Update parser.py

parser.py

@@ -33,7 +33,7 @@ def get_category(node, parent=None):
         if parent and isinstance(parent, ast.Return):
             return 'returned_variable'
     else:
-        return 'other'
+        return 'other'  # Default to 'other' for unrecognized nodes
 
 def create_vector(category, level, location, total_lines, parent_path):
     """Create a 6D vector optimized for role similarity, integrating variable roles into category_id."""
@@ -43,7 +43,7 @@ def create_vector(category, level, location, total_lines, parent_path):
         'other': 11, 'elif': 12, 'else': 13, 'except': 14, 'finally': 15, 'return': 16,
         'assigned_variable': 17, 'input_variable': 18, 'returned_variable': 19
     }
-    category_id = category_map.get(category, 0)
+    category_id = category_map.get(category, 0)  # Default to 0 for unknown categories
     start_line, end_line = location
     span = (end_line - start_line + 1) / total_lines
     center_pos = ((start_line + end_line) / 2) / total_lines
@@ -75,7 +75,10 @@ def parse_node(node, lines, prev_end, level=0, total_lines=None, parent_path=Non
     if any(line in processed_lines for line in range(start_line, end_line + 1)):
         return parts
 
-    category = get_category(node)
+    # Get category, default to 'other' if None
+    category = get_category(node, parent_path[-1] if parent_path else None) or 'other'
+    if category not in counters:
+        category = 'other'  # Ensure category exists in counters
     counters[category] += 1
     node_id = f"{category.capitalize()}[{counters[category]}]"
 
@@ -118,14 +121,17 @@ def parse_node(node, lines, prev_end, level=0, total_lines=None, parent_path=Non
         for arg in node.args.args:
             var_start = start_line  # Assume args are on the same line as function def for simplicity
             if var_start not in processed_lines:
-                counters['input_variable'] += 1
-                var_node_id = f"InputVariable[{counters['input_variable']}]"
+                arg_category = get_category(arg, node) or 'input_variable'
+                if arg_category not in counters:
+                    arg_category = 'input_variable'
+                counters[arg_category] += 1
+                var_node_id = f"InputVariable[{counters[arg_category]}]"
                 parts.append({
-                    'category': 'input_variable',
+                    'category': arg_category,
                     'source': f"    {arg.arg},",  # Indented as part of function
                     'location': (var_start, var_start),
                     'level': level + 1,
-                    'vector': create_vector('input_variable', level + 1, (var_start, var_start), total_lines, current_path),
+                    'vector': create_vector(arg_category, level + 1, (var_start, var_start), total_lines, current_path),
                     'parent_path': f"{current_path[0]} -> {var_node_id}",
                     'node_id': var_node_id
                 })