Update parser.py

parser.py

@@ -2,15 +2,11 @@
 import ast
 
 def get_category(node):
-    """Determine the category of an AST node."""
+    """Determine the category of an AST node or variable context."""
     if isinstance(node, (ast.Import, ast.ImportFrom)):
         return 'import'
-    elif isinstance(node, (ast.Assign, ast.AnnAssign, ast.AugAssign)):
-        return 'assignment'
-    elif isinstance(node, ast.FunctionDef):
+    elif isinstance(node, ast.FunctionDef) or isinstance(node, ast.AsyncFunctionDef):
         return 'function'
-    elif isinstance(node, ast.AsyncFunctionDef):
-        return 'async_function'
     elif isinstance(node, ast.ClassDef):
         return 'class'
     elif isinstance(node, ast.If):
@@ -27,15 +23,30 @@ def get_category(node):
         return 'expression'
     elif isinstance(node, ast.ExceptHandler):
         return 'except'
+    elif isinstance(node, ast.Assign) or isinstance(node, ast.AnnAssign) or isinstance(node, ast.AugAssign):
+        return 'assigned_variable'
     else:
         return 'other'
 
-def create_vector(category, level, location, total_lines, parent_path):
-    """Create a vector optimized for role similarity."""
+def get_variable_role(node, parent):
+    """Determine the role of a variable (input, assigned, returned) based on context."""
+    if isinstance(parent, (ast.FunctionDef, ast.AsyncFunctionDef)):
+        if isinstance(node, ast.arg):
+            return 'input_variable'
+        elif isinstance(parent, ast.Return) and isinstance(node, ast.Name):
+            return 'returned_variable'
+    elif isinstance(parent, (ast.Assign, ast.AnnAssign, ast.AugAssign)):
+        if isinstance(node, ast.Name) and node.id in [target.id for target in parent.targets if isinstance(target, ast.Name)]:
+            return 'assigned_variable'
+    return None
+
+def create_vector(category, level, location, total_lines, parent_path, variable_info=None):
+    """Create a vector optimized for role similarity, including variable info if applicable."""
     category_map = {
-        'import': 1, 'assignment': 2, 'function': 3, 'async_function': 4, 'class': 5,
-        'if': 6, 'while': 7, 'for': 8, 'try': 9, 'expression': 10, 'spacer': 11,
-        'other': 12, 'elif': 13, 'else': 14, 'except': 15, 'finally': 16, 'return': 17
+        'import': 1, 'function': 2, 'async_function': 3, 'class': 4,
+        'if': 5, 'while': 6, 'for': 7, 'try': 8, 'expression': 9, 'spacer': 10,
+        '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)
     start_line, end_line = location
@@ -44,7 +55,16 @@ def create_vector(category, level, location, total_lines, parent_path):
     parent_depth = len(parent_path)
     parent_weight = sum(category_map.get(parent.split('[')[0].lower(), 0) * (1 / (i + 1)) 
                         for i, parent in enumerate(parent_path)) / max(1, len(category_map))
-    return [category_id, level, center_pos, span, parent_depth, parent_weight]
+    
+    # Extend vector with variable-specific info if applicable
+    variable_role_id = 0
+    variable_name = ''
+    if variable_info:
+        role_map = {'input_variable': 1, 'assigned_variable': 2, 'returned_variable': 3}
+        variable_role_id = role_map.get(variable_info['role'], 0)
+        variable_name = variable_info['name'][:10]  # Truncate for brevity, hash if needed
+
+    return [category_id, level, center_pos, span, parent_depth, parent_weight, variable_role_id, hash(variable_name) % 1000 / 1000]
 
 def is_blank_or_comment(line):
     """Check if a line is blank or a comment."""
@@ -57,7 +77,7 @@ def parse_node(node, lines, prev_end, level=0, total_lines=None, parent_path=Non
     if parent_path is None:
         parent_path = []
     if counters is None:
-        counters = {cat: 0 for cat in ['import', 'assignment', 'function', 'async_function', 'class', 'if', 'while', 'for', 'try', 'return', 'expression', 'other', 'spacer', 'elif', 'else', 'except', 'finally']}
+        counters = {cat: 0 for cat in ['import', 'function', 'async_function', 'class', 'if', 'while', 'for', 'try', 'return', 'expression', 'other', 'spacer', 'elif', 'else', 'except', 'finally', 'assigned_variable', 'input_variable', 'returned_variable']}
     if processed_lines is None:
         processed_lines = set()
 
@@ -107,6 +127,25 @@ def parse_node(node, lines, prev_end, level=0, total_lines=None, parent_path=Non
         })
         processed_lines.add(start_line)
 
+    # Handle variables in function definitions (input variables)
+    if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.args.args:
+        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']}]"
+                var_info = {'role': 'input_variable', 'name': arg.arg}
+                parts.append({
+                    'category': 'input_variable',
+                    '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, var_info),
+                    'parent_path': f"{current_path[0]} -> {var_node_id}",
+                    'node_id': var_node_id
+                })
+                processed_lines.add(var_start)
+
     # Process nested bodies
     nested_prev_end = start_line
     for attr in ('body', 'orelse', 'handlers', 'finalbody'):
@@ -168,6 +207,43 @@ def parse_node(node, lines, prev_end, level=0, total_lines=None, parent_path=Non
                         parts.extend(child_parts)
                         nested_prev_end = max(nested_prev_end, child_parts[-1]['location'][1] if child_parts else child_start)
                     else:
+                        # Handle assignments and returns for variable detection
+                        if isinstance(child, ast.Assign) or isinstance(child, ast.AnnAssign) or isinstance(child, ast.AugAssign):
+                            for target in child.targets:
+                                if isinstance(target, ast.Name):
+                                    var_start = child.lineno
+                                    if var_start not in processed_lines and not is_blank_or_comment(lines[var_start - 1]):
+                                        counters['assigned_variable'] += 1
+                                        var_node_id = f"AssignedVariable[{counters['assigned_variable']}]"
+                                        var_info = {'role': 'assigned_variable', 'name': target.id}
+                                        parts.append({
+                                            'category': 'assigned_variable',
+                                            'source': lines[var_start - 1],
+                                            'location': (var_start, var_start),
+                                            'level': level + 1,
+                                            'vector': create_vector('assigned_variable', level + 1, (var_start, var_start), total_lines, current_path, var_info),
+                                            'parent_path': f"{current_path[0]} -> {var_node_id}",
+                                            'node_id': var_node_id
+                                        })
+                                        processed_lines.add(var_start)
+                        elif isinstance(child, ast.Return):
+                            for value in ast.walk(child):
+                                if isinstance(value, ast.Name):
+                                    var_start = child.lineno
+                                    if var_start not in processed_lines and not is_blank_or_comment(lines[var_start - 1]):
+                                        counters['returned_variable'] += 1
+                                        var_node_id = f"ReturnedVariable[{counters['returned_variable']}]"
+                                        var_info = {'role': 'returned_variable', 'name': value.id}
+                                        parts.append({
+                                            'category': 'returned_variable',
+                                            'source': lines[var_start - 1],
+                                            'location': (var_start, var_start),
+                                            'level': level + 1,
+                                            'vector': create_vector('returned_variable', level + 1, (var_start, var_start), total_lines, current_path, var_info),
+                                            'parent_path': f"{current_path[0]} -> {var_node_id}",
+                                            'node_id': var_node_id
+                                        })
+                                        processed_lines.add(var_start)
                         child_parts = parse_node(child, lines, nested_prev_end, level + 1, total_lines, current_path, counters, processed_lines)
                         parts.extend(child_parts)
                         nested_prev_end = child_parts[-1]['location'][1] if child_parts else nested_prev_end