Update parser.py

parser.py

@@ -73,12 +73,12 @@ def parse_node(node, lines, prev_end, level=0, total_lines=None, parent_path=Non
 
     # Skip if any lines are already processed
     if any(line in processed_lines for line in range(start_line, end_line + 1)):
-        return parts
+        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'  # Ensure category exists in counters
+        category = 'other'
     counters[category] += 1
     node_id = f"{category.capitalize()}[{counters[category]}]"
 
@@ -117,6 +117,7 @@ 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)
+    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
@@ -136,6 +137,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
@@ -160,8 +162,10 @@ 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)
-                        child_parts = parse_node(child, lines, child_start, level + 1, total_lines, current_path, counters, processed_lines)
+                            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]):
@@ -177,8 +181,10 @@ 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)
-                        child_parts = parse_node(child, lines, child_start, level + 1, total_lines, current_path, counters, processed_lines)
+                            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]):
@@ -194,8 +200,10 @@ 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)
-                        child_parts = parse_node(child, lines, child_start, level + 1, total_lines, current_path, counters, processed_lines)
+                            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
@@ -218,6 +226,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('assigned_variable')
                             else:  # AnnAssign or AugAssign
                                 target = child.target
                                 if isinstance(target, ast.Name):
@@ -235,6 +244,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('assigned_variable')
                         elif isinstance(child, ast.Return):
                             for value in ast.walk(child):
                                 if isinstance(value, ast.Name):
@@ -252,8 +262,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)
-                        child_parts = parse_node(child, lines, nested_prev_end, level + 1, total_lines, current_path, counters, processed_lines)
+                                        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
@@ -265,7 +277,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
+    return parts, category_sequence
 
 def parse_python_code(code):
     lines = code.splitlines(keepends=True)
@@ -273,15 +285,17 @@ 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], 'parent_path': 'Top-Level', 'node_id': 'Error[1]'}], ['error'])
 
     parts = []
     prev_end = 0
     processed_lines = set()
+    category_sequence = []
 
     for stmt in tree.body:
-        stmt_parts = parse_node(stmt, lines, prev_end, total_lines=total_lines, processed_lines=processed_lines)
+        stmt_parts, stmt_seq = parse_node(stmt, lines, prev_end, total_lines=total_lines, processed_lines=processed_lines)
         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:
@@ -304,8 +318,9 @@ def parse_python_code(code):
                             'node_id': spacer_node_id
                         })
                         processed_lines.add(i)
+                        category_sequence.append('spacer')
 
-    return parts
+    return parts, category_sequence
 
 def is_blank_or_comment(line):
     """Check if a line is blank or a comment."""