Update chunk_python_code.py

chunk_python_code.py

@@ -1,6 +1,9 @@
 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.
@@ -19,13 +22,13 @@ def chunk_python_code_with_metadata(source_code, reference):
     elif reference.startswith("doc/"):
         usage = "doc"
     else:
-        usage = "undefined" 
+        usage = "undefined"
 
     # Add metadata for usage to all documents
     for doc in documents:
         doc.metadata["reference"] = reference
         doc.metadata["usage"] = usage  # Add the determined usage metadata
-
+        print(doc)
     return documents
 
 
@@ -47,12 +50,10 @@ def iterate_ast(source_code, documents, reference):
         handle_no_first_level_node_found(documents, source_code, imports_dict, 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)
     
-
     # Iterate over first-level nodes
     for first_level_node in ast.iter_child_nodes(tree):
         if isinstance(first_level_node, ast.ClassDef):
@@ -62,22 +63,23 @@ def iterate_ast(source_code, documents, reference):
         elif isinstance(first_level_node, ast.Assign):
             handle_first_level_assign(first_level_node, documents, source_code, imports_dict)
 
-            
 
-def  handle_first_level_imports_only(documents, source_code, imports_dict, reference):
-        # Check if the file path before ".py" is "__init__"
+def handle_first_level_imports_only(documents, source_code, imports_dict, reference):
+    """
+    Handles cases where the first-level nodes are only imports.
+    """
     if reference.endswith("__init__.py"):
         type = "__init__-file"
     else:
         type = "undefined"
 
+    # Create metadata with "none" if imports are empty
+    metadata = {"type": type, "imports": "none" if not imports_dict else imports_dict}
+
     # Create and store a Document with the full source code
     doc = Document(
         page_content=source_code,
-        metadata={
-            "type": type,
-            "imports": imports_dict
-        }
+        metadata=metadata
     )
     documents.append(doc)
 
@@ -99,6 +101,7 @@ def extract_imports(tree):
                 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.
@@ -109,173 +112,136 @@ def analyze_imports(node, imports_dict):
             relevant_imports.add(imports_dict[sub_node.id])
     return list(relevant_imports)
 
-def handle_not_yet_defined_first_level_cases(documents, source_code, imports_dict):
-    if source_code: 
-        doc = Document(
-            page_content=source_code,
-            metadata={
-                "type": "undefined",
-                "imports": imports_dict
-            }
-        )
-        documents.append(doc)
-
-
 
 def handle_no_first_level_node_found(documents, source_code, imports_dict, reference):
     """
     Handles cases where no top-level nodes are found in the AST.
-    Stores the full content (likely comments) in a Document object
-    with metadata indicating type 'no code' or 'init' based on the reference.
     """
-    # Check if the file path before ".py" is "__init__"
     if reference.endswith("__init__.py"):
         type = "__init__-file"
     else:
         type = "undefined"
 
+    # Create metadata with "none" if imports are empty
+    metadata = {"type": type, "imports": "none" if not imports_dict else imports_dict}
+
     # Create and store a Document with the full source code
     doc = Document(
         page_content=source_code,
-        metadata={
-            "type": type,
-            "imports": imports_dict
-        }
+        metadata=metadata
     )
     documents.append(doc)
 
 
 def handle_first_level_assign(assign_node, documents, source_code, imports_dict):
     """
-    Handles assignment statements at the first level of the AST by storing them
-    in a Document object with metadata, including relevant imports.
+    Handles assignment statements at the first level of the AST.
     """
-    # Extract assignment source code
     assign_start_line = assign_node.lineno
     assign_end_line = assign_node.end_lineno
     assign_source = '\n'.join(source_code.splitlines()[assign_start_line-1:assign_end_line])
 
-    # Extract relevant imports for this assignment
+    # 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 and store Document for the assignment
     doc = Document(
         page_content=assign_source,
-        metadata={
-            "type": "Assign",
-            "imports": assign_imports
-        }
+        metadata=metadata
     )
     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, imports_dict):
     """
-    Handles classes at the first level of the AST by storing them
-    in a Document object with metadata, including relevant imports.
+    Handles classes at the first level of the AST.
     """
-    # Extract class source code
     class_start_line = class_node.lineno
-
-    # Find the line where the first function (def) starts or the next top-level node
     class_body_lines = [child.lineno for child in class_node.body if isinstance(child, ast.FunctionDef)]
-    class_end_line = min(class_body_lines, default=class_node.end_lineno) - 1  # Use `-1` to exclude the next node
-
-    # Generate the class source code
+    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 for this class
+
+    # Extract relevant imports
     class_imports = analyze_imports(class_node, imports_dict)
-    
+
+    # Create metadata with "none" if imports are empty
+    metadata = {
+        "type": "class",
+        "class": class_node.name,
+        "visibility": "public",
+        "imports": "none" if not class_imports else class_imports
+    }
+
     # Create and store Document for the class
     doc = Document(
         page_content=class_source,
-        metadata={
-            "type": "class",
-            "class": class_node.name,
-            "visibility": "public",
-            "imports": class_imports,  
-        }
+        metadata=metadata
     )
     documents.append(doc)
 
-
     # Handle methods within the class
     for second_level_node in ast.iter_child_nodes(class_node):
         if isinstance(second_level_node, ast.FunctionDef):
-            # Extract method source code
             method_start_line = (
                 second_level_node.decorator_list[0].lineno
                 if second_level_node.decorator_list else second_level_node.lineno
             )
             method_end_line = second_level_node.end_lineno
             method_source = '\n'.join(source_code.splitlines()[method_start_line-1:method_end_line])
-            
-            # Determine visibility metadata
+
             visibility = "internal" if second_level_node.name.startswith("_") else "public"
-            # Extract relevant imports for this method
             method_imports = analyze_imports(second_level_node, imports_dict)
-            
-            # Create and store Document
 
             doc = Document(
                 page_content=method_source,
                 metadata={
                     "type": "method",
                     "method": second_level_node.name,
-                    "visibility": "visibility",
-                    "imports": method_imports,  
+                    "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):
     """
-    Handles functions at the first level of the AST by storing them
-    in a Document object with metadata, including relevant imports.
+    Handles functions at the first level of the AST.
     """
-    # Extract function source code
     function_start_line = (
         function_node.decorator_list[0].lineno
         if function_node.decorator_list else function_node.lineno
     )
     function_end_line = function_node.end_lineno
     function_source = '\n'.join(source_code.splitlines()[function_start_line-1:function_end_line])
-    
-    # Determine visibility metadata
+
+    # Extract relevant imports
+    function_imports = analyze_imports(function_node, imports_dict)
+
     visibility = "internal" if function_node.name.startswith("_") else "public"
-    
-    # Check if the function is a CLI command (e.g., decorated with `@apy_command`)
+
     is_command = any(
-        decorator.id == "apy_command"  # Check decorator name
+        decorator.id == "apy_command"
         for decorator in function_node.decorator_list
-        if hasattr(decorator, "id")  # Ensure the decorator has an identifier
+        if hasattr(decorator, "id")
     )
-    
-    # Extract relevant imports for this function
-    function_imports = analyze_imports(function_node, imports_dict)
-    
-    # Create and store Document
-    
+
+    metadata = {
+        "type": "command" if is_command else "function",
+        "visibility": visibility,
+        "imports": "none" if not function_imports else function_imports
+    }
     if is_command:
-        doc = Document(
-            page_content=function_source,
-            metadata={
-                "type": "command",
-                "command": function_node.name,
-                "visibility": "public",
-                "imports": function_imports
-            }
-        )
+        metadata["command"] = function_node.name
     else:
-          doc = Document(
-            page_content=function_source,
-            metadata={
-                "type": "function",
-                "method": function_node.name,
-                "visibility": visibility,
-                "imports": function_imports
-            }
-        )          
-    documents.append(doc)
+        metadata["method"] = function_node.name
+
+    doc = Document(
+        page_content=function_source,
+        metadata=metadata
+    )
+    documents.append(doc)