Update parser.py

parser.py

@@ -13,13 +13,38 @@ def get_category(node):
         return 'async_function'
     elif isinstance(node, ast.ClassDef):
         return 'class'
+    elif isinstance(node, ast.If):
+        return 'if'
+    elif isinstance(node, ast.While):
+        return 'while'
+    elif isinstance(node, ast.For):
+        return 'for'
     elif isinstance(node, ast.Expr):
         return 'expression'
     else:
         return 'other'
 
-def parse_node(node, lines, prev_end, level=0):
+def create_vector(category, level, location, total_lines):
+    """Create a vector representation for a code part."""
+    # Vector: [category_id, level, start_line_normalized, end_line_normalized]
+    category_map = {
+        'import': 1, 'assignment': 2, 'function': 3, 'async_function': 4, 'class': 5,
+        'if': 6, 'while': 7, 'for': 8, 'expression': 9, 'spacer': 10, 'other': 11
+    }
+    category_id = category_map.get(category, 0)
+    start_line, end_line = location
+    return [
+        category_id,
+        level,
+        start_line / total_lines,  # Normalized start position
+        end_line / total_lines     # Normalized end position
+    ]
+
+def parse_node(node, lines, prev_end, level=0, total_lines=None):
     """Recursively parse an AST node and its children, assigning hierarchy levels."""
+    if total_lines is None:
+        total_lines = len(lines)
+
     parts = []
     start_line = getattr(node, 'lineno', prev_end + 1)
     end_line = getattr(node, 'end_lineno', start_line)
@@ -27,56 +52,71 @@ def parse_node(node, lines, prev_end, level=0):
     # Handle spacers before the node
     if start_line > prev_end + 1:
         spacer_lines = lines[prev_end:start_line - 1]
+        spacer_vector = create_vector('spacer', level, (prev_end + 1, start_line - 1), total_lines)
         parts.append({
             'category': 'spacer',
             'source': ''.join(spacer_lines),
             'location': (prev_end + 1, start_line - 1),
-            'level': level
+            'level': level,
+            'vector': spacer_vector
         })
 
     # Capture the node's source
     stmt_lines = lines[start_line - 1:end_line]
+    node_vector = create_vector(get_category(node), level, (start_line, end_line), total_lines)
     parts.append({
         'category': get_category(node),
         'source': ''.join(stmt_lines),
         'location': (start_line, end_line),
-        'level': level
+        'level': level,
+        'vector': node_vector
     })
 
-    # Process nested nodes (e.g., class or function bodies)
+    # Process nested nodes (e.g., class/function bodies, control structures)
     if hasattr(node, 'body'):
         nested_prev_end = end_line - 1
         for child in node.body:
-            child_parts = parse_node(child, lines, nested_prev_end, level + 1)
+            child_parts = parse_node(child, lines, nested_prev_end, level + 1, total_lines)
             parts.extend(child_parts)
             nested_prev_end = child_parts[-1]['location'][1]
 
+    # Handle additional bodies (e.g., else, elif, orelse for loops)
+    if hasattr(node, 'orelse') and node.orelse:
+        orelse_prev_end = parts[-1]['location'][1]
+        for child in node.orelse:
+            child_parts = parse_node(child, lines, orelse_prev_end, level + 1, total_lines)
+            parts.extend(child_parts)
+            orelse_prev_end = child_parts[-1]['location'][1]
+
     return parts
 
 def parse_python_code(code):
-    """Parse Python code string and return parts with hierarchy."""
+    """Parse Python code string and return parts with hierarchy and vectors."""
     lines = code.splitlines(keepends=True)
+    total_lines = len(lines)
     try:
         tree = ast.parse(code)
     except SyntaxError:
-        return [{'category': 'error', 'source': 'Invalid Python code', 'location': (1, 1), 'level': 0}]
+        return [{'category': 'error', 'source': 'Invalid Python code', 'location': (1, 1), 'level': 0, 'vector': [0, 0, 1.0, 1.0]}]
 
     parts = []
     prev_end = 0
 
     for stmt in tree.body:
-        stmt_parts = parse_node(stmt, lines, prev_end)
+        stmt_parts = parse_node(stmt, lines, prev_end, total_lines=total_lines)
         parts.extend(stmt_parts)
         prev_end = stmt_parts[-1]['location'][1]
 
     # Capture trailing spacers
-    if prev_end < len(lines):
+    if prev_end < total_lines:
         remaining_lines = lines[prev_end:]
+        spacer_vector = create_vector('spacer', 0, (prev_end + 1, total_lines + 1), total_lines)
         parts.append({
             'category': 'spacer',
             'source': ''.join(remaining_lines),
-            'location': (prev_end + 1, len(lines) + 1),
-            'level': 0
+            'location': (prev_end + 1, total_lines + 1),
+            'level': 0,
+            'vector': spacer_vector
         })
 
     return parts