Update app.py

app.py

@@ -34,7 +34,7 @@ def parse_code():
 
     parts, _ = parse_python_code(code)
     
-    # Extract nodes (functions and variables) and connections
+    # Extract nodes and connections
     nodes = []
     connections = []
     node_id_map = {}  # Maps parser node_id to canvas node index
@@ -45,7 +45,7 @@ def parse_code():
         node_id = part['node_id']
         source = part['source'].strip()
 
-        if category in ['function', 'assigned_variable', 'input_variable', 'returned_variable']:
+        if category in ['function', 'assigned_variable', 'input_variable', 'returned_variable', 'import']:
             node_data = {
                 'id': len(nodes),
                 'type': category,
@@ -58,31 +58,34 @@ def parse_code():
             }
 
             if category == 'function':
-                # Extract function name and parameters
+                # Function: inputs are parameters, output is return
                 if 'def ' in source:
                     func_name = source.split('def ')[1].split('(')[0]
                     node_data['label'] = func_name
                     params = source.split('(')[1].split(')')[0].split(',')
                     params = [p.strip() for p in params if p.strip()]
                     node_data['inputs'] = params
-                    # Assume function can have outputs (return statements)
                     node_data['outputs'] = ['return']
             elif category == 'input_variable':
-                # Input variable (function parameter)
+                # Input variable: output to parent function
                 var_name = source.strip().rstrip(',')
                 node_data['label'] = var_name
                 node_data['outputs'] = [var_name]
             elif category == 'assigned_variable':
-                # Assigned variable
+                # Assigned variable: input is value, output is variable
                 var_name = source.split('=')[0].strip()
                 node_data['label'] = var_name
                 node_data['inputs'] = ['value']
                 node_data['outputs'] = [var_name]
             elif category == 'returned_variable':
-                # Returned variable
+                # Returned variable: input from function
                 var_name = source.split('return ')[1].strip() if 'return ' in source else node_id
                 node_data['label'] = var_name
                 node_data['inputs'] = [var_name]
+            elif category == 'import':
+                # Import: no inputs/outputs, just a node
+                import_name = source.split('import ')[1].split()[0] if 'import ' in source else node_id
+                node_data['label'] = import_name
 
             nodes.append(node_data)
             node_id_map[node_id] = node_data['id']
@@ -90,17 +93,29 @@ def parse_code():
 
     # Create connections based on parent_path and variable usage
     for part in parts:
-        if part['category'] in ['input_variable', 'assigned_variable', 'returned_variable']:
+        category = part['category']
+        if category in ['input_variable', 'returned_variable', 'assigned_variable']:
             parent_path = part['parent_path']
             if 'Function' in parent_path:
                 parent_node_id = parent_path.split(' -> ')[0]
                 if parent_node_id in node_id_map:
                     from_id = node_id_map[part['node_id']]
                     to_id = node_id_map[parent_node_id]
-                    if part['category'] == 'input_variable':
+                    if category == 'input_variable':
                         connections.append({'from': from_id, 'to': to_id})
-                    elif part['category'] == 'returned_variable':
+                    elif category == 'returned_variable':
                         connections.append({'from': to_id, 'to': from_id})
+                    elif category == 'assigned_variable':
+                        # Connect assigned variables to parent function (e.g., local variables)
+                        connections.append({'from': from_id, 'to': to_id})
+        elif category == 'import':
+            # Connect imports to functions that might use them (simplified: connect to all functions)
+            for other_part in parts:
+                if other_part['category'] == 'function' and other_part['node_id'] in node_id_map:
+                    connections.append({
+                        'from': node_id_map[part['node_id']],
+                        'to': node_id_map[other_part['node_id']]
+                    })
 
     return jsonify({'nodes': nodes, 'connections': connections})