Update parser.py

parser.py

@@ -21,6 +21,8 @@ def get_category(node):
         return 'for'
     elif isinstance(node, ast.Try):
         return 'try'
+    elif isinstance(node, ast.Return):
+        return 'return'
     elif isinstance(node, ast.Expr):
         return 'expression'
     else:
@@ -31,29 +33,19 @@ def create_vector(category, level, location, total_lines, parent_path):
     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
+        'other': 12, 'elif': 13, 'else': 14, 'except': 15, 'finally': 16, 'return': 17
     }
     category_id = category_map.get(category, 0)
     start_line, end_line = location
-    span = (end_line - start_line + 1) / total_lines  # Normalized size of the part
-    center_pos = ((start_line + end_line) / 2) / total_lines  # Center position normalized
+    span = (end_line - start_line + 1) / total_lines
+    center_pos = ((start_line + end_line) / 2) / total_lines
     parent_depth = len(parent_path)
-    
-    # Weighted sum of parent categories (simple weighting by position)
     parent_weight = 0
     for i, parent in enumerate(parent_path):
         parent_category = parent.split('[')[0].lower()
-        parent_weight += category_map.get(parent_category, 0) * (1 / (i + 1))  # Decay with depth
-    parent_weight = parent_weight / max(1, len(category_map))  # Normalize by max category ID
-
-    return [
-        category_id,       # Type of the part
-        level,             # Nesting depth
-        center_pos,        # Center position in file (0.0 to 1.0)
-        span,              # Relative size in file (0.0 to 1.0)
-        parent_depth,      # Number of ancestors
-        parent_weight      # Semantic connection to parents (0.0 to 1.0-ish)
-    ]
+        parent_weight += category_map.get(parent_category, 0) * (1 / (i + 1))
+    parent_weight = parent_weight / max(1, len(category_map))
+    return [category_id, level, center_pos, span, parent_depth, parent_weight]
 
 def parse_node(node, lines, prev_end, level=0, total_lines=None, parent_path=None, counters=None):
     """Recursively parse an AST node with full hierarchy tracking."""
@@ -62,20 +54,22 @@ 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 = {'if': 0, 'while': 0, 'for': 0, 'function': 0, 'class': 0, 'try': 0}
+        counters = {cat: 0 for cat in ['import', 'assignment', 'function', 'async_function', 'class', 'if', 'while', 'for', 'try', 'return', 'expression', 'other', 'spacer', 'elif', 'else', 'except', 'finally']}
 
     parts = []
     start_line = getattr(node, 'lineno', prev_end + 1)
     end_line = getattr(node, 'end_lineno', start_line)
     category = get_category(node)
 
-    node_id = ''
-    if category in counters:
-        counters[category] += 1
-        node_id = f"{category.capitalize()}[{counters[category]}]"
+    # Assign a node_id to every part
+    counters[category] += 1
+    node_id = f"{category.capitalize()}[{counters[category]}]"
 
+    # Spacer before the node
     if start_line > prev_end + 1:
         spacer_lines = lines[prev_end:start_line - 1]
+        counters['spacer'] += 1
+        spacer_node_id = f"Spacer[{counters['spacer']}]"
         spacer_vector = create_vector('spacer', level, (prev_end + 1, start_line - 1), total_lines, parent_path)
         parts.append({
             'category': 'spacer',
@@ -84,11 +78,12 @@ def parse_node(node, lines, prev_end, level=0, total_lines=None, parent_path=Non
             'level': level,
             'vector': spacer_vector,
             'parent_path': ' -> '.join(parent_path) if parent_path else 'Top-Level',
-            'node_id': ''
+            'node_id': spacer_node_id
         })
 
+    # Main node
     stmt_lines = lines[start_line - 1:end_line]
-    current_path = parent_path + ([node_id] if node_id else [])
+    current_path = parent_path + [node_id]
     node_vector = create_vector(category, level, (start_line, end_line), total_lines, current_path)
     parts.append({
         'category': category,
@@ -100,6 +95,7 @@ def parse_node(node, lines, prev_end, level=0, total_lines=None, parent_path=Non
         'node_id': node_id
     })
 
+    # Process nested bodies
     nested_prev_end = end_line - 1
     for attr in ('body', 'orelse', 'handlers', 'finalbody'):
         if hasattr(node, attr) and getattr(node, attr):
@@ -107,6 +103,8 @@ def parse_node(node, lines, prev_end, level=0, total_lines=None, parent_path=Non
             for child in getattr(node, attr):
                 if attr == 'orelse' and isinstance(node, ast.If) and child.lineno != start_line:
                     sub_category = 'elif' if child.lineno != end_line else 'else'
+                    counters[sub_category] += 1
+                    sub_node_id = f"{sub_category.capitalize()}[{counters[sub_category]}]"
                     sub_vector = create_vector(sub_category, level, (child.lineno, getattr(child, 'end_lineno', child.lineno)), total_lines, current_path)
                     sub_parts.append({
                         'category': sub_category,
@@ -115,11 +113,13 @@ def parse_node(node, lines, prev_end, level=0, total_lines=None, parent_path=Non
                         'level': level,
                         'vector': sub_vector,
                         'parent_path': ' -> '.join(parent_path) if parent_path else 'Top-Level',
-                        'node_id': node_id
+                        'node_id': node_id  # Inherits parent's node_id for context
                     })
                     child_parts = parse_node(child, lines, child.lineno - 1, level + 1, total_lines, current_path, counters)
                     sub_parts.extend(child_parts)
                 elif attr == 'handlers' and isinstance(child, ast.ExceptHandler):
+                    counters['except'] += 1
+                    sub_node_id = f"Except[{counters['except']}]"
                     sub_vector = create_vector('except', level, (child.lineno, getattr(child, 'end_lineno', child.lineno)), total_lines, current_path)
                     sub_parts.append({
                         'category': 'except',
@@ -133,6 +133,8 @@ def parse_node(node, lines, prev_end, level=0, total_lines=None, parent_path=Non
                     child_parts = parse_node(child, lines, child.lineno - 1, level + 1, total_lines, current_path, counters)
                     sub_parts.extend(child_parts)
                 elif attr == 'finalbody':
+                    counters['finally'] += 1
+                    sub_node_id = f"Finally[{counters['finally']}]"
                     sub_vector = create_vector('finally', level, (child.lineno, getattr(child, 'end_lineno', child.lineno)), total_lines, current_path)
                     sub_parts.append({
                         'category': 'finally',
@@ -160,7 +162,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': ''}]
+        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
@@ -172,6 +174,9 @@ def parse_python_code(code):
 
     if prev_end < total_lines:
         remaining_lines = lines[prev_end:]
+        counters = {'spacer': 0}
+        counters['spacer'] += 1
+        spacer_node_id = f"Spacer[{counters['spacer']}]"
         spacer_vector = create_vector('spacer', 0, (prev_end + 1, total_lines + 1), total_lines, [])
         parts.append({
             'category': 'spacer',
@@ -180,7 +185,7 @@ def parse_python_code(code):
             'level': 0,
             'vector': spacer_vector,
             'parent_path': 'Top-Level',
-            'node_id': ''
+            'node_id': spacer_node_id
         })
 
     return parts