Update parser.py

parser.py

@@ -1,8 +1,8 @@
 # parser.py
 import ast
 
-def get_category(node):
-    """Determine the category of an AST node or variable context."""
+def get_category(node, parent=None):
+    """Determine the category of an AST node or variable context, including variable roles."""
     if isinstance(node, (ast.Import, ast.ImportFrom)):
         return 'import'
     elif isinstance(node, ast.FunctionDef) or isinstance(node, ast.AsyncFunctionDef):
@@ -24,32 +24,19 @@ def get_category(node):
     elif isinstance(node, ast.ExceptHandler):
         return 'except'
     elif isinstance(node, (ast.Assign, ast.AnnAssign, ast.AugAssign)):
-        return 'assigned_variable'
-    else:
-        return 'other'
-
-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):
+        if parent and isinstance(parent, (ast.FunctionDef, ast.AsyncFunctionDef, ast.If, ast.Try, ast.While, ast.For)):
+            return 'assigned_variable'
+    elif isinstance(node, ast.arg):  # Input variables in function definitions
+        if parent and isinstance(parent, (ast.FunctionDef, ast.AsyncFunctionDef)):
             return 'input_variable'
-        elif isinstance(parent, ast.Return) and isinstance(node, ast.Name):
+    elif isinstance(node, ast.Name):  # Returned variables in return statements
+        if parent and isinstance(parent, ast.Return):
             return 'returned_variable'
-    elif isinstance(parent, (ast.Assign, ast.AnnAssign, ast.AugAssign)):
-        if isinstance(node, ast.Name):
-            # Handle different target structures
-            if isinstance(parent, ast.Assign):
-                for target in parent.targets:
-                    if isinstance(target, ast.Name) and target.id == node.id:
-                        return 'assigned_variable'
-            elif isinstance(parent, (ast.AnnAssign, ast.AugAssign)):
-                target = parent.target
-                if isinstance(target, ast.Name) and target.id == node.id:
-                    return 'assigned_variable'
-    return None
+    else:
+        return 'other'
 
-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."""
+def create_vector(category, level, location, total_lines, parent_path):
+    """Create a 6D vector optimized for role similarity, integrating variable roles into category_id."""
     category_map = {
         'import': 1, 'function': 2, 'async_function': 3, 'class': 4,
         'if': 5, 'while': 6, 'for': 7, 'try': 8, 'expression': 9, 'spacer': 10,
@@ -63,16 +50,7 @@ def create_vector(category, level, location, total_lines, parent_path, variable_
     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))
-    
-    # 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]
+    return [category_id, level, center_pos, span, parent_depth, parent_weight]
 
 def is_blank_or_comment(line):
     """Check if a line is blank or a comment."""
@@ -142,13 +120,12 @@ def parse_node(node, lines, prev_end, level=0, total_lines=None, parent_path=Non
             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),
+                    'vector': create_vector('input_variable', level + 1, (var_start, var_start), total_lines, current_path),
                     'parent_path': f"{current_path[0]} -> {var_node_id}",
                     'node_id': var_node_id
                 })
@@ -216,42 +193,42 @@ def parse_node(node, lines, prev_end, level=0, total_lines=None, parent_path=Non
                         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):
-                            for target in child.targets:
+                        if isinstance(child, (ast.Assign, ast.AnnAssign, ast.AugAssign)):
+                            # Handle different target structures
+                            if isinstance(child, ast.Assign):
+                                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']}]"
+                                            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),
+                                                'parent_path': f"{current_path[0]} -> {var_node_id}",
+                                                'node_id': var_node_id
+                                            })
+                                            processed_lines.add(var_start)
+                            else:  # AnnAssign or AugAssign
+                                target = child.target
                                 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),
+                                            'vector': create_vector('assigned_variable', level + 1, (var_start, var_start), total_lines, current_path),
                                             'parent_path': f"{current_path[0]} -> {var_node_id}",
                                             'node_id': var_node_id
                                         })
                                         processed_lines.add(var_start)
-                        elif isinstance(child, ast.AnnAssign) or isinstance(child, ast.AugAssign):
-                            target = child.target
-                            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):
@@ -259,13 +236,12 @@ def parse_node(node, lines, prev_end, level=0, total_lines=None, parent_path=Non
                                     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),
+                                            'vector': create_vector('returned_variable', level + 1, (var_start, var_start), total_lines, current_path),
                                             'parent_path': f"{current_path[0]} -> {var_node_id}",
                                             'node_id': var_node_id
                                         })
@@ -291,7 +267,7 @@ def parse_python_code(code):
     try:
         tree = ast.parse(code)
     except SyntaxError:
-        return [{'category': 'error', 'source': 'Invalid Python code', 'location': (1, 1), 'level': 0, 'vector': [0, 0, 1.0, 0.0, 0, 0, 0, 0], 'parent_path': 'Top-Level', 'node_id': 'Error[1]'}]
+        return [{'category': 'error', 'source': 'Invalid Python code', 'location': (1, 1), 'level': 0, 'vector': [0, 0, 1.0, 0.0, 0, 0], 'parent_path': 'Top-Level', 'node_id': 'Error[1]'}]
 
     parts = []
     prev_end = 0