Update parser.py

parser.py

@@ -1,6 +1,4 @@
 # parser.py
-import os
-os.system("pip install --upgrade ast")
 import ast
 
 def get_category(node):
@@ -25,7 +23,7 @@ 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):
+    elif isinstance(node, (ast.Assign, ast.AnnAssign, ast.AugAssign)):
         return 'assigned_variable'
     else:
         return 'other'
@@ -38,8 +36,16 @@ def get_variable_role(node, parent):
         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'
+        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
 
 def create_vector(category, level, location, total_lines, parent_path, variable_info=None):
@@ -130,25 +136,23 @@ 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(child, ast.Assign) or isinstance(child, ast.AnnAssign) or isinstance(child, ast.AugAssign):
-        targets = child.targets if isinstance(child, ast.Assign) else [child.target]
-        for target in 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)
+    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
@@ -212,11 +216,8 @@ 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) or isinstance(child, ast.AnnAssign) or isinstance(child, ast.AugAssign):
-                            targets = child.target
-                            if isinstance(child, ast.Assign):
-                                targets = child.targets
-                            for target in targets:
+                        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]):
@@ -233,6 +234,24 @@ def parse_node(node, lines, prev_end, level=0, total_lines=None, parent_path=Non
                                             '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):
@@ -272,7 +291,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], '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, 0, 0], 'parent_path': 'Top-Level', 'node_id': 'Error[1]'}]
 
     parts = []
     prev_end = 0