Update parser.py

parser.py

@@ -2,9 +2,7 @@
 import ast
 
 def get_category(node):
-    """
-    Determine the category of an AST node based on its type.
-    """
+    """Determine the category of an AST node."""
     if isinstance(node, (ast.Import, ast.ImportFrom)):
         return 'import'
     elif isinstance(node, (ast.Assign, ast.AnnAssign, ast.AugAssign)):
@@ -20,52 +18,65 @@ def get_category(node):
     else:
         return 'other'
 
-def parse_python_file(file_path):
-    """
-    Parse a Python file and return a list of its parts with categories, source code, and locations.
+def parse_node(node, lines, prev_end, level=0):
+    """Recursively parse an AST node and its children, assigning hierarchy levels."""
+    parts = []
+    start_line = getattr(node, 'lineno', prev_end + 1)
+    end_line = getattr(node, 'end_lineno', start_line)
+
+    # Handle spacers before the node
+    if start_line > prev_end + 1:
+        spacer_lines = lines[prev_end:start_line - 1]
+        parts.append({
+            'category': 'spacer',
+            'source': ''.join(spacer_lines),
+            'location': (prev_end + 1, start_line - 1),
+            'level': level
+        })
+
+    # Capture the node's source
+    stmt_lines = lines[start_line - 1:end_line]
+    parts.append({
+        'category': get_category(node),
+        'source': ''.join(stmt_lines),
+        'location': (start_line, end_line),
+        'level': level
+    })
 
-    Args:
-        file_path (str): Path to the Python file to parse.
+    # Process nested nodes (e.g., class or function bodies)
+    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)
+            parts.extend(child_parts)
+            nested_prev_end = child_parts[-1]['location'][1]
 
-    Returns:
-        list: A list of dictionaries, each containing 'category', 'source', and 'location'.
-    """
-    with open(file_path, 'r') as file:
-        code = file.read()
+    return parts
 
+def parse_python_code(code):
+    """Parse Python code string and return parts with hierarchy."""
     lines = code.splitlines(keepends=True)
-    tree = ast.parse(code)
+    try:
+        tree = ast.parse(code)
+    except SyntaxError:
+        return [{'category': 'error', 'source': 'Invalid Python code', 'location': (1, 1), 'level': 0}]
 
     parts = []
     prev_end = 0
 
     for stmt in tree.body:
-        start_line = stmt.lineno
-        end_line = getattr(stmt, 'end_lineno', start_line)
-
-        if start_line > prev_end + 1:
-            spacer_lines = lines[prev_end:start_line - 1]
-            parts.append({
-                'category': 'spacer',
-                'source': ''.join(spacer_lines),
-                'location': (prev_end + 1, start_line - 1)
-            })
-
-        stmt_lines = lines[start_line - 1:end_line]
-        parts.append({
-            'category': get_category(stmt),
-            'source': ''.join(stmt_lines),
-            'location': (start_line, end_line)
-        })
-
-        prev_end = end_line
+        stmt_parts = parse_node(stmt, lines, prev_end)
+        parts.extend(stmt_parts)
+        prev_end = stmt_parts[-1]['location'][1]
 
+    # Capture trailing spacers
     if prev_end < len(lines):
         remaining_lines = lines[prev_end:]
         parts.append({
             'category': 'spacer',
             'source': ''.join(remaining_lines),
-            'location': (prev_end + 1, len(lines) + 1)
+            'location': (prev_end + 1, len(lines) + 1),
+            'level': 0
         })
 
     return parts