Update parser.py

parser.py

@@ -1,332 +1,134 @@
 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, ast.AsyncFunctionDef)):
-        return '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.Try):
-        return 'try'
-    elif isinstance(node, ast.Return):
-        return 'return'
-    elif isinstance(node, ast.Expr):
-        return 'expression'
-    elif isinstance(node, ast.ExceptHandler):
-        return 'except'
-    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):
-        if parent and isinstance(parent, (ast.FunctionDef, ast.AsyncFunctionDef)):
-            return 'input_variable'
-    elif isinstance(node, ast.Name):
-        if parent and isinstance(parent, ast.Return):
-            return 'returned_variable'
-    return 'other'
-
-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('#')
+import hashlib
+
+def get_category_id(category):
+    """Maps categorical roles to integers for vector embedding."""
+    mapping = {
+        'unknown': 0, 'import': 1, 'function': 2, 'class': 3, 
+        'if': 4, 'while': 5, 'for': 6, 'try': 7, 'expression': 8, 
+        'spacer': 9, 'elif': 10, 'else': 11, 'except': 12, 
+        'return': 13, 'assigned_variable': 14, 'variable_def': 15
+    }
+    return mapping.get(category, 0)
 
 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
+    """
+    Creates a 6D normalized vector:
+    [Category, Depth, RelativeCenter, Density, ParentDepth, AncestryWeight]
+    """
+    cat_id = get_category_id(category)
+    start, end = location
+    total_lines = max(1, total_lines)
+    
+    # metrics
+    span = (end - start + 1) / total_lines
+    center = ((start + end) / 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 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:
-        total_lines = len(lines)
-    if parent_path is None:
-        parent_path = []
-    if counters is None:
-        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()
-
-    parts = []
-    start_line = getattr(node, 'lineno', prev_end + 1)
-    end_line = getattr(node, 'end_lineno', start_line)
-
-    if any(line in processed_lines for line in range(start_line, end_line + 1)):
-        return parts, []
-
-    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
-
-    if start_line > prev_end + 1:
-        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_path = parent_path + [node_id]
-    if start_line not in processed_lines and not is_blank_or_comment(lines[start_line - 1]):
-        part = {
-            'category': category,
-            'source': lines[start_line - 1],
-            'location': (start_line, start_line),
-            'level': level,
-            '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)
-
-    if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.args.args:
-        for arg in node.args.args:
-            var_start = start_line
-            if var_start not in processed_lines:
-                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},",
-                    'location': (var_start, var_start),
-                    'level': level + 1,
-                    'vector': create_vector(arg_category, 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)
-
-    nested_prev_end = start_line
-    for attr in ('body', 'orelse', 'handlers', 'finalbody'):
-        if hasattr(node, attr) and getattr(node, attr):
-            for child in getattr(node, attr):
-                child_start = getattr(child, 'lineno', nested_prev_end + 1)
-                child_end = getattr(child, 'end_lineno', child_start)
-                if not any(line in processed_lines for line in range(child_start, child_end + 1)):
-                    if attr == 'orelse' and isinstance(node, ast.If) and child_start != start_line:
-                        sub_category = 'elif' if 'elif' in lines[child_start - 1] else 'else'
-                        if child_start not in processed_lines and not is_blank_or_comment(lines[child_start - 1]):
-                            counters[sub_category] += 1
-                            sub_node_id = f"{sub_category.capitalize()}[{counters[sub_category]}]"
-                            parts.append({
-                                'category': sub_category,
-                                'source': lines[child_start - 1],
-                                'location': (child_start, child_start),
-                                'level': level,
-                                'vector': create_vector(sub_category, level, (child_start, child_start), total_lines, current_path),
-                                'parent_path': ' -> '.join(parent_path) if parent_path else 'Top-Level',
-                                'node_id': sub_node_id
-                            })
-                            processed_lines.add(child_start)
-                        child_parts, child_seq = parse_node(child, lines, child_start, level + 1, total_lines, current_path, counters, processed_lines)
-                        parts.extend(child_parts)
-                        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]):
-                            counters['except'] += 1
-                            sub_node_id = f"Except[{counters['except']}]"
-                            parts.append({
-                                'category': 'except',
-                                'source': lines[child_start - 1],
-                                'location': (child_start, child_start),
-                                'level': level,
-                                'vector': create_vector('except', level, (child_start, child_start), total_lines, current_path),
-                                'parent_path': ' -> '.join(parent_path) if parent_path else 'Top-Level',
-                                'node_id': sub_node_id
-                            })
-                            processed_lines.add(child_start)
-                        child_parts, child_seq = parse_node(child, lines, child_start, level + 1, total_lines, current_path, counters, processed_lines)
-                        parts.extend(child_parts)
-                        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]):
-                            counters['finally'] += 1
-                            sub_node_id = f"Finally[{counters['finally']}]"
-                            parts.append({
-                                'category': 'finally',
-                                'source': lines[child_start - 1],
-                                'location': (child_start, child_start),
-                                'level': level,
-                                'vector': create_vector('finally', level, (child_start, child_start), total_lines, current_path),
-                                'parent_path': ' -> '.join(parent_path) if parent_path else 'Top-Level',
-                                'node_id': sub_node_id
-                            })
-                            processed_lines.add(child_start)
-                        child_parts, child_seq = parse_node(child, lines, child_start, level + 1, total_lines, current_path, counters, processed_lines)
-                        parts.extend(child_parts)
-                        nested_prev_end = max(nested_prev_end, child_parts[-1]['location'][1] if child_parts else child_start)
-                    else:
-                        if isinstance(child, (ast.Assign, ast.AnnAssign, ast.AugAssign)):
-                            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],
-                                            '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,
-                                            'value': value
-                                        })
-                                        processed_lines.add(var_start)
-                        elif isinstance(child, ast.Return):
-                            for value in ast.walk(child):
-                                if isinstance(value, ast.Name):
-                                    var_start = child.lineno
-                                    if var_start not in processed_lines and not is_blank_or_comment(lines[var_start - 1]):
-                                        counters['returned_variable'] += 1
-                                        var_node_id = f"ReturnedVariable[{counters['returned_variable']}]"
-                                        parts.append({
-                                            'category': 'returned_variable',
-                                            'source': lines[var_start - 1],
-                                            'location': (var_start, var_start),
-                                            'level': level + 1,
-                                            'vector': create_vector('returned_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)
-                        child_parts, child_seq = parse_node(child, lines, nested_prev_end, level + 1, total_lines, current_path, counters, processed_lines)
-                        parts.extend(child_parts)
-                        nested_prev_end = child_parts[-1]['location'][1] if child_parts else nested_prev_end
-
-    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:
-            parts[-1]['location'] = (start_line, final_end)
-            parts[-1]['source'] = ''.join(lines[start_line - 1:final_end])
-            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, []
-
-def parse_python_code(code):
-    lines = code.splitlines(keepends=True)
-    total_lines = len(lines)
+    
+    # Ancestry weight: Simple hash sum of parent IDs to represent unique path
+    path_str = "".join(parent_path)
+    parent_weight = (int(hashlib.md5(path_str.encode()).hexdigest(), 16) % 100) / 100.0
+    
+    return [
+        cat_id, 
+        level, 
+        float(f"{center:.4f}"), 
+        float(f"{span:.4f}"), 
+        parent_depth, 
+        float(f"{parent_weight:.4f}")
+    ]
+
+def parse_source_to_graph(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]'}], ['error'])
+    except SyntaxError as e:
+        return {"error": f"Syntax Error on line {e.lineno}: {e.msg}"}
 
-    parts = []
-    prev_end = 0
-    processed_lines = set()
-
-    var_defs, var_uses = collect_variable_usage(tree)
-
-    for stmt in tree.body:
-        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)
-        prev_end = stmt_parts[-1]['location'][1] if stmt_parts else prev_end
+    lines = code.splitlines(keepends=True)
+    total_lines = len(lines)
+    nodes = []
+    
+    # Recursive visitor
+    def traverse(node, parent_path=[], level=0, parent_id=None):
+        category = 'other'
+        name = getattr(node, 'name', None)
+        # Unique Node ID based on position to ensure consistency
+        node_id = f"{type(node).__name__}_{getattr(node, 'lineno', 0)}_{getattr(node, 'col_offset', 0)}"
+        
+        # Categorization logic
+        if isinstance(node, (ast.Import, ast.ImportFrom)): category = 'import'; name = "import"
+        elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): category = 'function'
+        elif isinstance(node, ast.ClassDef): category = 'class'
+        elif isinstance(node, ast.If): category = 'if'; name = "if"
+        elif isinstance(node, (ast.For, ast.AsyncFor)): category = 'for'; name = "for"
+        elif isinstance(node, ast.While): category = 'while'; name = "while"
+        elif isinstance(node, ast.Return): category = 'return'; name = "return"
+        elif isinstance(node, (ast.Assign, ast.AnnAssign)): category = 'assigned_variable'; name = "assignment"
+        elif isinstance(node, ast.Expr): category = 'expression'; name = "expr"
+        elif isinstance(node, ast.Try): category = 'try'; name = "try"
+        elif isinstance(node, ast.ExceptHandler): category = 'except'; name = "except"
+        
+        lineno = getattr(node, 'lineno', 0)
+        end_lineno = getattr(node, 'end_lineno', lineno)
+        
+        if lineno == 0: return # Skip nodes without line numbers (e.g. Load context)
+
+        # Create source snippet
+        source_segment = "".join(lines[lineno-1:end_lineno])
+        
+        # Determine Label
+        label = name if name else category
+        if category == 'assigned_variable':
+            targets = getattr(node, 'targets', []) or [getattr(node, 'target', None)]
+            if targets and isinstance(targets[0], ast.Name):
+                label = f"{targets[0].id} ="
+        
+        vector = create_vector(category, level, (lineno, end_lineno), total_lines, parent_path)
+        
+        node_data = {
+            "id": node_id,
+            "label": label,
+            "type": category,
+            "source": source_segment.strip(),
+            "vector": vector,
+            "level": level,
+            "lineno": lineno,
+            "parent_id": parent_id
+        }
+        
+        # Filter: Only visualize structural elements (skip raw expressions unless useful)
+        if category != 'other':
+            nodes.append(node_data)
+            current_path = parent_path + [node_id]
+            current_parent = node_id
+            next_level = level + 1
+        else:
+            current_path = parent_path
+            current_parent = parent_id
+            next_level = level
 
-    if prev_end < total_lines:
-        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)
+        for child in ast.iter_child_nodes(node):
+            traverse(child, current_path, next_level, current_parent)
 
-    return parts, []
+    for node in tree.body:
+        traverse(node)
+    
+    # Sort by line number for linear visual flow
+    nodes.sort(key=lambda x: x['lineno'])
+    
+    return {"nodes": nodes, "connections": generate_connections(nodes)}
+
+def generate_connections(nodes):
+    connections = []
+    node_map = {n['id']: n for n in nodes}
+    
+    for node in nodes:
+        # 1. Structural Hierarchy (Tree)
+        if node['parent_id'] and node['parent_id'] in node_map:
+            connections.append({
+                "from": node['parent_id'], 
+                "to": node['id'], 
+                "type": "hierarchy"
+            })
+            
+    return connections