✨ feat(code_parser.py): add support for parsing return statement in function or method nodes

🐛 fix(code_parser.py): fix import formatting to improve readability
✨ feat(code_parser.py): add support for parsing function arguments and body in function or method nodes
This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-12-01 13:12:19 -03:00
commit 4583d30f4c

View file

@ -6,7 +6,8 @@ from typing import Any, Dict, List, Type, Union
from cachetools import TTLCache, cachedmethod, keys from cachetools import TTLCache, cachedmethod, keys
from fastapi import HTTPException from fastapi import HTTPException
from langflow.interface.custom.schema import CallableCodeDetails, ClassCodeDetails from langflow.interface.custom.schema import (CallableCodeDetails,
ClassCodeDetails)
class CodeSyntaxError(HTTPException): class CodeSyntaxError(HTTPException):
@ -56,6 +57,9 @@ class CodeParser:
ast.Assign: self.parse_global_vars, ast.Assign: self.parse_global_vars,
} }
def __get_tree(self): def __get_tree(self):
""" """
Parses the provided code to validate its syntax. Parses the provided code to validate its syntax.
@ -79,6 +83,7 @@ class CodeParser:
if handler := self.handlers.get(type(node)): # type: ignore if handler := self.handlers.get(type(node)): # type: ignore
handler(node) # type: ignore handler(node) # type: ignore
def parse_imports(self, node: Union[ast.Import, ast.ImportFrom]) -> None: def parse_imports(self, node: Union[ast.Import, ast.ImportFrom]) -> None:
""" """
Extracts "imports" from the code, including aliases. Extracts "imports" from the code, including aliases.
@ -149,12 +154,16 @@ class CodeParser:
# Handle cases where the type is not found in the constructed environment # Handle cases where the type is not found in the constructed environment
pass pass
func = CallableCodeDetails( func = CallableCodeDetails(
name=node.name, doc=ast.get_docstring(node), args=[], body=[], return_type=return_type or get_data_type() name=node.name,
doc=ast.get_docstring(node),
args= self.parse_function_args(node),
body= self.parse_function_body(node),
return_type=return_type or get_data_type(),
has_return=self.parse_return_statement(node),
) )
func.args = self.parse_function_args(node)
func.body = self.parse_function_body(node)
return func.model_dump() return func.model_dump()
@ -230,6 +239,14 @@ class CodeParser:
""" """
return [ast.unparse(line) for line in node.body] return [ast.unparse(line) for line in node.body]
def parse_return_statement(self, node: ast.FunctionDef) -> bool:
"""
Parses the return statement of a function or method node.
"""
return any(isinstance(n, ast.Return) for n in node.body)
def parse_assign(self, stmt): def parse_assign(self, stmt):
""" """
Parses an Assign statement and returns a dictionary Parses an Assign statement and returns a dictionary