Update chunk_python_code.py

chunk_python_code.py

@@ -66,79 +66,6 @@ def _iterate_ast(python_code, documents, file_path):
                 _handle_not_defined_case(python_code))
              
 
-def _chunk_import_only_python_code(python_code, file_path):
-    """
-    Handles cases where the first-level nodes are only imports.
-    """
-    documents = []
-    if file_path.endswith("__init__.py"):
-        type = "__init__-file"
-    else:
-        type = "undefined"
-
-    # Create metadata without imports
-    metadata = {"type": type}
-
-    # Create and store a Document with the full source code
-    doc = Document(
-        page_content=python_code,
-        metadata=metadata
-    )
-    documents.append(doc)
-    return documents
-
-
-
-def _handle_not_defined_case(python_code):
-    documents = []
-    documents.extend(
-        _chunk_python_code_by_character(python_code)  
-    return documents
-    
-
-def _chunk_nodeless_python_code(python_code, file_path):
-    """
-    Handles cases where no top-level nodes are found in the AST.
-    """
-    documents = []
-    if file_path.endswith("__init__.py"):
-        type = "__init__-file"
-    else:
-        type = "undefined"
-
-    # Create metadata without imports
-    metadata = {"type": type}
-
-    # Create and store a Document with the full source code
-    doc = Document(
-        page_content=python_code,
-        metadata=metadata
-    )
-    documents.append(doc)
-
-    return documents
-
-def _chunk_first_level_assign_node(ast_node, python_code):
-    
-    """
-    Handles assignment statements at the first level of the AST.
-    """
-    documents = []
-    assign_start_line = ast_node.lineno
-    assign_end_line = ast_node.end_lineno
-    assign_source = '\n'.join(python_code.splitlines()[assign_start_line-1:assign_end_line])
-
-    # Create metadata without imports
-    metadata = {"type": "Assign"}
-
-    # Create and store Document for the assignment
-    doc = Document(
-        page_content=assign_source,
-        metadata=metadata
-    )
-    documents.append(doc)
-
-    return documents
 
 def _handle_first_level_class(ast_node , python_code):
     """
@@ -188,6 +115,14 @@ def _handle_first_level_class(ast_node , python_code):
 
     return documents
 
+
+def _handle_not_defined_case(python_code):
+    documents = []
+    documents.extend(
+        _chunk_python_code_by_character(python_code))
+    return documents
+    
+
 def _chunk_first_level_func_node(ast_node, python_code):
     """
     Handles functions at the first level of the AST.
@@ -226,6 +161,76 @@ def _chunk_first_level_func_node(ast_node, python_code):
     return documents
 
 
+
+def _chunk_first_level_assign_node(ast_node, python_code):
+    
+    """
+    Handles assignment statements at the first level of the AST.
+    """
+    documents = []
+    assign_start_line = ast_node.lineno
+    assign_end_line = ast_node.end_lineno
+    assign_source = '\n'.join(python_code.splitlines()[assign_start_line-1:assign_end_line])
+
+    # Create metadata without imports
+    metadata = {"type": "Assign"}
+
+    # Create and store Document for the assignment
+    doc = Document(
+        page_content=assign_source,
+        metadata=metadata
+    )
+    documents.append(doc)
+
+    return documents
+
+
+    
+def _chunk_import_only_python_code(python_code, file_path):
+    """
+    Handles cases where the first-level nodes are only imports.
+    """
+    documents = []
+    if file_path.endswith("__init__.py"):
+        type = "__init__-file"
+    else:
+        type = "undefined"
+
+    # Create metadata without imports
+    metadata = {"type": type}
+
+    # Create and store a Document with the full source code
+    doc = Document(
+        page_content=python_code,
+        metadata=metadata
+    )
+    documents.append(doc)
+    return documents
+
+def _chunk_nodeless_python_code(python_code, file_path):
+    """
+    Handles cases where no top-level nodes are found in the AST.
+    """
+    documents = []
+    if file_path.endswith("__init__.py"):
+        type = "__init__-file"
+    else:
+        type = "undefined"
+
+    # Create metadata without imports
+    metadata = {"type": type}
+
+    # Create and store a Document with the full source code
+    doc = Document(
+        page_content=python_code,
+        metadata=metadata
+    )
+    documents.append(doc)
+
+    return documents
+
+
+
 from langchain.text_splitter import RecursiveCharacterTextSplitter