Update parser.py

parser.py

@@ -1,11 +1,11 @@
-# parser.py
 import ast
+from collections import defaultdict
 
 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):
+    elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
         return 'function'
     elif isinstance(node, ast.ClassDef):
         return 'class'
@@ -26,36 +26,63 @@ def get_category(node, parent=None):
     elif isinstance(node, (ast.Assign, ast.AnnAssign, ast.AugAssign)):
         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
+    elif isinstance(node, ast.arg):
         if parent and isinstance(parent, (ast.FunctionDef, ast.AsyncFunctionDef)):
             return 'input_variable'
-    elif isinstance(node, ast.Name):  # Returned variables in return statements
+    elif isinstance(node, ast.Name):
         if parent and isinstance(parent, ast.Return):
             return 'returned_variable'
-    else:
-        return 'other'  # Default to 'other' for unrecognized nodes
+    return 'other'
 
-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,
-        '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)  # 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
-    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]
+def get_value(node):
+    """Extract the value of an AST node (e.g., for assignment)."""
+    if isinstance(node, ast.Constant):
+        return str(node.value)
+    elif isinstance(node, ast.Name):
+        return node.id
+    elif isinstance(node, ast.BinOp):
+        return '<expression>'
+    elif isinstance(node, ast.Call):
+        return '<function_call>'
+    return '<complex>'
 
-def is_blank_or_comment(line):
-    """Check if a line is blank or a comment."""
-    stripped = line.strip()
-    return not stripped or stripped.startswith('#')
+def collect_variable_usage(tree):
+    """Collect definitions and uses of variables, respecting scope."""
+    var_defs = defaultdict(list)  # {var_name: [(node_id, scope)]}
+    var_uses = defaultdict(list)  # {var_name: [(node_id, scope)]}
+    scope_stack = ['global']
+
+    def traverse(node, current_scope):
+        node_id = getattr(node, 'node_id', None)
+        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
+            scope_stack.append(node.name)
+            current_scope = node.name
+            for arg in node.args.args:
+                var_defs[arg.arg].append((f"InputVariable[{arg.arg}]", current_scope))
+            for body_node in node.body:
+                traverse(body_node, current_scope)
+            scope_stack.pop()
+        elif isinstance(node, (ast.Assign, ast.AnnAssign)):
+            value = get_value(node.value) if hasattr(node, 'value') else '<unknown>'
+            for target in (node.targets if isinstance(node, ast.Assign) else [node.target]):
+                if isinstance(target, ast.Name):
+                    var_defs[target.id].append((node_id, current_scope))
+                    # Parse value for uses
+                    for child in ast.walk(node.value):
+                        if isinstance(child, ast.Name):
+                            var_uses[child.id].append((node_id, current_scope))
+        elif isinstance(node, ast.Name) and isinstance(node.ctx, ast.Load):
+            var_uses[node.id].append((node_id, current_scope))
+        for child in ast.iter_child_nodes(node):
+            traverse(child, current_scope)
+
+    for node in tree.body:
+        node_id = getattr(node, 'node_id', None)
+        if node_id:
+            setattr(node, 'node_id', node_id)
+        traverse(node, 'global')
+
+    return var_defs, var_uses
 
 def parse_node(node, lines, prev_end, level=0, total_lines=None, parent_path=None, counters=None, processed_lines=None):
     if total_lines is None:
@@ -63,7 +90,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', 'function', 'async_function', 'class', 'if', 'while', 'for', 'try', 'return', 'expression', 'other', 'spacer', 'elif', 'else', 'except', 'finally', 'assigned_variable', 'input_variable', 'returned_variable']}
+        counters = {cat: 0 for cat in ['import', '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()
 
@@ -71,41 +98,35 @@ def parse_node(node, lines, prev_end, level=0, total_lines=None, parent_path=Non
     start_line = getattr(node, 'lineno', prev_end + 1)
     end_line = getattr(node, 'end_lineno', start_line)
 
-    # Skip if any lines are already processed
     if any(line in processed_lines for line in range(start_line, end_line + 1)):
         return parts, []
 
-    # 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'
     counters[category] += 1
     node_id = f"{category.capitalize()}[{counters[category]}]"
+    setattr(node, 'node_id', node_id)  # Attach node_id to AST node
 
-    # Spacer before node (only for blank lines or comments)
     if start_line > prev_end + 1:
-        spacer_lines = lines[prev_end:start_line - 1]
-        spacer_lines_set = set(range(prev_end + 1, start_line))
-        if not spacer_lines_set.issubset(processed_lines):
-            for i, line in enumerate(spacer_lines, prev_end + 1):
-                if i not in processed_lines and is_blank_or_comment(line):
-                    counters['spacer'] += 1
-                    spacer_node_id = f"Spacer[{counters['spacer']}]"
-                    parts.append({
-                        'category': 'spacer',
-                        'source': line,
-                        'location': (i, i),
-                        'level': level,
-                        'vector': create_vector('spacer', level, (i, i), total_lines, parent_path),
-                        'parent_path': ' -> '.join(parent_path) if parent_path else 'Top-Level',
-                        'node_id': spacer_node_id
-                    })
-                    processed_lines.add(i)
+        for i, line in enumerate(lines[prev_end:start_line - 1], prev_end + 1):
+            if i not in processed_lines and is_blank_or_comment(line):
+                counters['spacer'] += 1
+                spacer_node_id = f"Spacer[{counters['spacer']}]"
+                parts.append({
+                    'category': 'spacer',
+                    'source': line,
+                    'location': (i, i),
+                    'level': level,
+                    'vector': create_vector('spacer', level, (i, i), total_lines, parent_path),
+                    'parent_path': ' -> '.join(parent_path) if parent_path else 'Top-Level',
+                    'node_id': spacer_node_id
+                })
+                processed_lines.add(i)
 
-    # Current node's header (e.g., 'def', 'if', 'try')
     current_path = parent_path + [node_id]
     if start_line not in processed_lines and not is_blank_or_comment(lines[start_line - 1]):
-        parts.append({
+        part = {
             'category': category,
             'source': lines[start_line - 1],
             'location': (start_line, start_line),
@@ -113,23 +134,22 @@ def parse_node(node, lines, prev_end, level=0, total_lines=None, parent_path=Non
             'vector': create_vector(category, level, (start_line, start_line), total_lines, current_path),
             'parent_path': ' -> '.join(parent_path) if parent_path else 'Top-Level',
             'node_id': node_id
-        })
+        }
+        if category == 'assigned_variable':
+            part['value'] = get_value(node.value) if hasattr(node, 'value') else '<unknown>'
+        parts.append(part)
         processed_lines.add(start_line)
 
-    # Handle variables in function definitions (input variables)
-    category_sequence = [category]
     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
+            var_start = start_line
             if var_start not in processed_lines:
-                arg_category = get_category(arg, node) or 'input_variable'
-                if arg_category not in counters:
-                    arg_category = 'input_variable'
+                arg_category = 'input_variable'
                 counters[arg_category] += 1
                 var_node_id = f"InputVariable[{counters[arg_category]}]"
                 parts.append({
                     'category': arg_category,
-                    'source': f"    {arg.arg},",  # Indented as part of function
+                    'source': f"    {arg.arg},",
                     'location': (var_start, var_start),
                     'level': level + 1,
                     'vector': create_vector(arg_category, level + 1, (var_start, var_start), total_lines, current_path),
@@ -137,9 +157,7 @@ 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)
-                category_sequence.append(arg_category)
 
-    # Process nested bodies
     nested_prev_end = start_line
     for attr in ('body', 'orelse', 'handlers', 'finalbody'):
         if hasattr(node, attr) and getattr(node, attr):
@@ -162,10 +180,8 @@ def parse_node(node, lines, prev_end, level=0, total_lines=None, parent_path=Non
                                 'node_id': sub_node_id
                             })
                             processed_lines.add(child_start)
-                            category_sequence.append(sub_category)
                         child_parts, child_seq = parse_node(child, lines, child_start, level + 1, total_lines, current_path, counters, processed_lines)
                         parts.extend(child_parts)
-                        category_sequence.extend(child_seq)
                         nested_prev_end = max(nested_prev_end, child_parts[-1]['location'][1] if child_parts else child_start)
                     elif attr == 'handlers' and isinstance(child, ast.ExceptHandler):
                         if child_start not in processed_lines and not is_blank_or_comment(lines[child_start - 1]):
@@ -181,10 +197,8 @@ def parse_node(node, lines, prev_end, level=0, total_lines=None, parent_path=Non
                                 'node_id': sub_node_id
                             })
                             processed_lines.add(child_start)
-                            category_sequence.append('except')
                         child_parts, child_seq = parse_node(child, lines, child_start, level + 1, total_lines, current_path, counters, processed_lines)
                         parts.extend(child_parts)
-                        category_sequence.extend(child_seq)
                         nested_prev_end = max(nested_prev_end, child_parts[-1]['location'][1] if child_parts else child_start)
                     elif attr == 'finalbody':
                         if child_start not in processed_lines and not is_blank_or_comment(lines[child_start - 1]):
@@ -200,40 +214,18 @@ def parse_node(node, lines, prev_end, level=0, total_lines=None, parent_path=Non
                                 'node_id': sub_node_id
                             })
                             processed_lines.add(child_start)
-                            category_sequence.append('finally')
                         child_parts, child_seq = parse_node(child, lines, child_start, level + 1, total_lines, current_path, counters, processed_lines)
                         parts.extend(child_parts)
-                        category_sequence.extend(child_seq)
                         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, 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)
-                                            category_sequence.append('assigned_variable')
-                            else:  # AnnAssign or AugAssign
-                                target = child.target
+                            for target in (child.targets if isinstance(child, ast.Assign) else [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']}]"
+                                        value = get_value(child.value) if hasattr(child, 'value') else '<unknown>'
                                         parts.append({
                                             'category': 'assigned_variable',
                                             'source': lines[var_start - 1],
@@ -241,10 +233,10 @@ def parse_node(node, lines, prev_end, level=0, total_lines=None, parent_path=Non
                                             '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
+                                            'node_id': var_node_id,
+                                            'value': value
                                         })
                                         processed_lines.add(var_start)
-                                        category_sequence.append('assigned_variable')
                         elif isinstance(child, ast.Return):
                             for value in ast.walk(child):
                                 if isinstance(value, ast.Name):
@@ -262,13 +254,10 @@ 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)
-                                        category_sequence.append('returned_variable')
                         child_parts, child_seq = parse_node(child, lines, nested_prev_end, level + 1, total_lines, current_path, counters, processed_lines)
                         parts.extend(child_parts)
-                        category_sequence.extend(child_seq)
                         nested_prev_end = child_parts[-1]['location'][1] if child_parts else nested_prev_end
 
-    # Update end_line and source of the parent node if its body extends it
     if nested_prev_end > start_line and start_line not in processed_lines:
         final_end = nested_prev_end
         if start_line not in processed_lines:
@@ -277,7 +266,7 @@ def parse_node(node, lines, prev_end, level=0, total_lines=None, parent_path=Non
             parts[-1]['vector'] = create_vector(category, level, (start_line, final_end), total_lines, current_path)
             processed_lines.update(range(start_line, final_end + 1))
 
-    return parts, category_sequence
+    return parts, []
 
 def parse_python_code(code):
     lines = code.splitlines(keepends=True)
@@ -290,39 +279,32 @@ def parse_python_code(code):
     parts = []
     prev_end = 0
     processed_lines = set()
-    category_sequence = []
+
+    var_defs, var_uses = collect_variable_usage(tree)
 
     for stmt in tree.body:
-        stmt_parts, stmt_seq = parse_node(stmt, lines, prev_end, total_lines=total_lines, processed_lines=processed_lines)
+        stmt_parts, _ = parse_node(stmt, lines, prev_end, total_lines=total_lines, processed_lines=processed_lines)
+        for part in stmt_parts:
+            part['var_defs'] = var_defs
+            part['var_uses'] = var_uses
         parts.extend(stmt_parts)
-        category_sequence.extend(stmt_seq)
         prev_end = stmt_parts[-1]['location'][1] if stmt_parts else prev_end
 
     if prev_end < total_lines:
-        remaining_lines = lines[prev_end:]
-        remaining_lines_set = set(range(prev_end + 1, total_lines + 1))
-        if not remaining_lines_set.issubset(processed_lines):
-            for i, line in enumerate(remaining_lines, prev_end + 1):
-                if i not in processed_lines:
-                    if is_blank_or_comment(line):
-                        counters = {'spacer': 0}
-                        counters['spacer'] += 1
-                        spacer_node_id = f"Spacer[{counters['spacer']}]"
-                        parts.append({
-                            'category': 'spacer',
-                            'source': line,
-                            'location': (i, i),
-                            'level': 0,
-                            'vector': create_vector('spacer', 0, (i, i), total_lines, []),
-                            'parent_path': 'Top-Level',
-                            'node_id': spacer_node_id
-                        })
-                        processed_lines.add(i)
-                        category_sequence.append('spacer')
-
-    return parts, category_sequence
+        for i, line in enumerate(lines[prev_end:], prev_end + 1):
+            if i not in processed_lines and is_blank_or_comment(line):
+                counters = {'spacer': 0}
+                counters['spacer'] += 1
+                spacer_node_id = f"Spacer[{counters['spacer']}]"
+                parts.append({
+                    'category': 'spacer',
+                    'source': line,
+                    'location': (i, i),
+                    'level': 0,
+                    'vector': create_vector('spacer', 0, (i, i), total_lines, []),
+                    'parent_path': 'Top-Level',
+                    'node_id': spacer_node_id
+                })
+                processed_lines.add(i)
 
-def is_blank_or_comment(line):
-    """Check if a line is blank or a comment."""
-    stripped = line.strip()
-    return not stripped or stripped.startswith('#')
+    return parts, []