Refactor code_parser.py and test_custom_component.py

src/backend/langflow/interface/custom/code_parser/code_parser.py

@@ -6,8 +6,8 @@ from typing import Any, Dict, List, Type, Union
 
 from cachetools import TTLCache, cachedmethod, keys
 from fastapi import HTTPException
-
+from langflow.interface.custom.schema import (CallableCodeDetails,
-from langflow.interface.custom.schema import CallableCodeDetails, ClassCodeDetails
+                                              ClassCodeDetails)
 
 
 class CodeSyntaxError(HTTPException):
@@ -57,7 +57,7 @@ class CodeParser:
             ast.Assign: self.parse_global_vars,
         }
 
-    def __get_tree(self):
+    def get_tree(self):
         """
         Parses the provided code to validate its syntax.
         It tries to parse the code into an abstract syntax tree (AST).
@@ -313,7 +313,7 @@ class CodeParser:
         """
         Runs all parsing operations and returns the resulting data.
         """
-        tree = self.__get_tree()
+        tree = self.get_tree()
 
         for node in ast.walk(tree):
             self.parse_node(node)

tests/test_custom_component.py

@@ -5,9 +5,10 @@ from uuid import uuid4
 import pytest
 from fastapi import HTTPException
 from langflow.interface.custom.base import CustomComponent
-from langflow.interface.custom.code_parser import CodeParser, CodeSyntaxError
+from langflow.interface.custom.code_parser.code_parser import (CodeParser,
-from langflow.interface.custom.component import (Component,
+                                                               CodeSyntaxError)
-                                                 ComponentCodeNullError)
+from langflow.interface.custom.custom_component.component import (
+    Component, ComponentCodeNullError)
 from langflow.interface.custom.utils import (build_custom_component_template,
                                              create_and_validate_component)
 from langflow.services.database.models.flow import Flow, FlowCreate
@@ -50,7 +51,7 @@ def test_code_parser_get_tree():
     Test the __get_tree method of the CodeParser class.
     """
     parser = CodeParser(code_default)
-    tree = parser._CodeParser__get_tree()
+    tree = parser.get_tree()
     assert isinstance(tree, ast.AST)
 
 
@@ -63,7 +64,7 @@ def test_code_parser_syntax_error():
 
     parser = CodeParser(code_syntax_error)
     with pytest.raises(CodeSyntaxError):
-        parser._CodeParser__get_tree()
+        parser.get_tree()
 
 
 def test_component_init():
@@ -140,7 +141,7 @@ def test_code_parser_parse_imports_import():
     class with an import statement.
     """
     parser = CodeParser(code_default)
-    tree = parser._CodeParser__get_tree()
+    tree = parser.get_tree()
     for node in ast.walk(tree):
         if isinstance(node, ast.Import):
             parser.parse_imports(node)
@@ -153,7 +154,7 @@ def test_code_parser_parse_imports_importfrom():
     class with an import from statement.
     """
     parser = CodeParser("from os import path")
-    tree = parser._CodeParser__get_tree()
+    tree = parser.get_tree()
     for node in ast.walk(tree):
         if isinstance(node, ast.ImportFrom):
             parser.parse_imports(node)
@@ -165,7 +166,7 @@ def test_code_parser_parse_functions():
     Test the parse_functions method of the CodeParser class.
     """
     parser = CodeParser("def test(): pass")
-    tree = parser._CodeParser__get_tree()
+    tree = parser.get_tree()
     for node in ast.walk(tree):
         if isinstance(node, ast.FunctionDef):
             parser.parse_functions(node)
@@ -178,7 +179,7 @@ def test_code_parser_parse_classes():
     Test the parse_classes method of the CodeParser class.
     """
     parser = CodeParser("class Test: pass")
-    tree = parser._CodeParser__get_tree()
+    tree = parser.get_tree()
     for node in ast.walk(tree):
         if isinstance(node, ast.ClassDef):
             parser.parse_classes(node)
@@ -191,7 +192,7 @@ def test_code_parser_parse_global_vars():
     Test the parse_global_vars method of the CodeParser class.
     """
     parser = CodeParser("x = 1")
-    tree = parser._CodeParser__get_tree()
+    tree = parser.get_tree()
     for node in ast.walk(tree):
         if isinstance(node, ast.Assign):
             parser.parse_global_vars(node)
@@ -312,7 +313,7 @@ def test_code_parser_parse_ann_assign():
     stmt = ast.AnnAssign(
         target=ast.Name(id="x", ctx=ast.Store()),
         annotation=ast.Name(id="int", ctx=ast.Load()),
-        value=ast.Num(n=1),
+        value=ast.Constant(n=1),
         simple=1,
     )
     result = parser.parse_ann_assign(stmt)