Update chunk_python_code.py

chunk_python_code.py

@@ -1,9 +1,6 @@
 import ast
 from langchain.schema import Document 
 
-import ast
-from langchain.schema import Document 
-
 def chunk_python_code_with_metadata(source_code, reference):
     """
     Entry point method to process the Python file.
@@ -40,31 +37,28 @@ def iterate_ast(source_code, documents, reference):
     # Parse the source code into an abstract syntax tree (AST)
     tree = ast.parse(source_code, filename=reference)
 
-    # Gather all top-level imports for later use
-    imports_dict = extract_imports(tree)
-
     first_level_nodes = list(ast.iter_child_nodes(tree))
 
     # Check if there are no first-level nodes
     if not first_level_nodes:
-        handle_no_first_level_node_found(documents, source_code, imports_dict, reference)
+        handle_no_first_level_node_found(documents, source_code, reference)
         return
 
     all_imports = all(isinstance(node, (ast.Import, ast.ImportFrom)) for node in first_level_nodes)
     if all_imports:
-        handle_first_level_imports_only(documents, source_code, imports_dict, reference)
+        handle_first_level_imports_only(documents, source_code, reference)
     
     # Iterate over first-level nodes
     for first_level_node in ast.iter_child_nodes(tree):
         if isinstance(first_level_node, ast.ClassDef):
-            handle_first_level_class(first_level_node, documents, source_code, imports_dict)
+            handle_first_level_class(first_level_node, documents, source_code)
         elif isinstance(first_level_node, ast.FunctionDef):
-            handle_first_level_func(first_level_node, documents, source_code, imports_dict)
+            handle_first_level_func(first_level_node, documents, source_code)
         elif isinstance(first_level_node, ast.Assign):
-            handle_first_level_assign(first_level_node, documents, source_code, imports_dict)
+            handle_first_level_assign(first_level_node, documents, source_code)
 
 
-def handle_first_level_imports_only(documents, source_code, imports_dict, reference):
+def handle_first_level_imports_only(documents, source_code, reference):
     """
     Handles cases where the first-level nodes are only imports.
     """
@@ -73,8 +67,8 @@ def handle_first_level_imports_only(documents, source_code, imports_dict, refere
     else:
         type = "undefined"
 
-    # Create metadata with "none" if imports are empty
-    metadata = {"type": type, "imports": "none" if not imports_dict else imports_dict}
+    # Create metadata without imports
+    metadata = {"type": type}
 
     # Create and store a Document with the full source code
     doc = Document(
@@ -84,36 +78,7 @@ def handle_first_level_imports_only(documents, source_code, imports_dict, refere
     documents.append(doc)
 
 
-def extract_imports(tree):
-    """
-    Extracts all import statements from the AST tree and organizes them
-    into a dictionary keyed by their fully qualified names for later analysis.
-    """
-    imports_dict = {}
-    for node in ast.walk(tree):
-        if isinstance(node, ast.Import):
-            for alias in node.names:
-                imports_dict[alias.name] = alias.name
-        elif isinstance(node, ast.ImportFrom):
-            module = node.module if node.module else ""
-            for alias in node.names:
-                full_name = f"{module}.{alias.name}" if module else alias.name
-                imports_dict[alias.name] = full_name
-    return imports_dict
-
-
-def analyze_imports(node, imports_dict):
-    """
-    Analyzes the given node's body and signature to find relevant imports.
-    """
-    relevant_imports = set()
-    for sub_node in ast.walk(node):
-        if isinstance(sub_node, ast.Name) and sub_node.id in imports_dict:
-            relevant_imports.add(imports_dict[sub_node.id])
-    return list(relevant_imports)
-
-
-def handle_no_first_level_node_found(documents, source_code, imports_dict, reference):
+def handle_no_first_level_node_found(documents, source_code, reference):
     """
     Handles cases where no top-level nodes are found in the AST.
     """
@@ -122,8 +87,8 @@ def handle_no_first_level_node_found(documents, source_code, imports_dict, refer
     else:
         type = "undefined"
 
-    # Create metadata with "none" if imports are empty
-    metadata = {"type": type, "imports": "none" if not imports_dict else imports_dict}
+    # Create metadata without imports
+    metadata = {"type": type}
 
     # Create and store a Document with the full source code
     doc = Document(
@@ -133,7 +98,7 @@ def handle_no_first_level_node_found(documents, source_code, imports_dict, refer
     documents.append(doc)
 
 
-def handle_first_level_assign(assign_node, documents, source_code, imports_dict):
+def handle_first_level_assign(assign_node, documents, source_code):
     """
     Handles assignment statements at the first level of the AST.
     """
@@ -141,11 +106,8 @@ def handle_first_level_assign(assign_node, documents, source_code, imports_dict)
     assign_end_line = assign_node.end_lineno
     assign_source = '\n'.join(source_code.splitlines()[assign_start_line-1:assign_end_line])
 
-    # Extract relevant imports
-    assign_imports = analyze_imports(assign_node, imports_dict)
-
-    # Create metadata with "none" if imports are empty
-    metadata = {"type": "Assign", "imports": "none" if not assign_imports else assign_imports}
+    # Create metadata without imports
+    metadata = {"type": "Assign"}
 
     # Create and store Document for the assignment
     doc = Document(
@@ -155,7 +117,7 @@ def handle_first_level_assign(assign_node, documents, source_code, imports_dict)
     documents.append(doc)
 
 
-def handle_first_level_class(class_node, documents, source_code, imports_dict):
+def handle_first_level_class(class_node, documents, source_code):
     """
     Handles classes at the first level of the AST.
     """
@@ -164,15 +126,11 @@ def handle_first_level_class(class_node, documents, source_code, imports_dict):
     class_end_line = min(class_body_lines, default=class_node.end_lineno) - 1
     class_source = '\n'.join(source_code.splitlines()[class_start_line-1:class_end_line])
 
-    # Extract relevant imports
-    class_imports = analyze_imports(class_node, imports_dict)
-
-    # Create metadata with "none" if imports are empty
+    # Create metadata without imports
     metadata = {
         "type": "class",
         "class": class_node.name,
-        "visibility": "public",
-        "imports": "none" if not class_imports else class_imports
+        "visibility": "public"
     }
 
     # Create and store Document for the class
@@ -193,7 +151,6 @@ def handle_first_level_class(class_node, documents, source_code, imports_dict):
             method_source = '\n'.join(source_code.splitlines()[method_start_line-1:method_end_line])
 
             visibility = "internal" if second_level_node.name.startswith("_") else "public"
-            method_imports = analyze_imports(second_level_node, imports_dict)
 
             doc = Document(
                 page_content=method_source,
@@ -201,14 +158,13 @@ def handle_first_level_class(class_node, documents, source_code, imports_dict):
                     "type": "method",
                     "method": second_level_node.name,
                     "visibility": visibility,
-                    "imports": "none" if not method_imports else method_imports,
                     "class": class_node.name
                 }
             )
             documents.append(doc)
 
 
-def handle_first_level_func(function_node, documents, source_code, imports_dict):
+def handle_first_level_func(function_node, documents, source_code):
     """
     Handles functions at the first level of the AST.
     """
@@ -219,9 +175,6 @@ def handle_first_level_func(function_node, documents, source_code, imports_dict)
     function_end_line = function_node.end_lineno
     function_source = '\n'.join(source_code.splitlines()[function_start_line-1:function_end_line])
 
-    # Extract relevant imports
-    function_imports = analyze_imports(function_node, imports_dict)
-
     visibility = "internal" if function_node.name.startswith("_") else "public"
 
     is_command = any(
@@ -232,8 +185,7 @@ def handle_first_level_func(function_node, documents, source_code, imports_dict)
 
     metadata = {
         "type": "command" if is_command else "function",
-        "visibility": visibility,
-        "imports": "none" if not function_imports else function_imports
+        "visibility": visibility
     }
     if is_command:
         metadata["command"] = function_node.name
@@ -245,3 +197,8 @@ def handle_first_level_func(function_node, documents, source_code, imports_dict)
         metadata=metadata
     )
     documents.append(doc)
+
+
+
+
+