Refactor construct_eval_env method in CodeParser

class
This commit is contained in:
Gabriel Luiz Freitas Almeida 2023-11-23 11:09:19 -03:00
commit 6fea589778

View file

@ -6,7 +6,6 @@ from typing import Any, Dict, List, Type, Union
from cachetools import TTLCache, cachedmethod from cachetools import TTLCache, cachedmethod
from fastapi import HTTPException from fastapi import HTTPException
from langflow.interface.custom.schema import CallableCodeDetails, ClassCodeDetails from langflow.interface.custom.schema import CallableCodeDetails, ClassCodeDetails
@ -105,13 +104,15 @@ class CodeParser:
arg_dict["type"] = ast.unparse(arg.annotation) arg_dict["type"] = ast.unparse(arg.annotation)
return arg_dict return arg_dict
def construct_eval_env(self, return_type_str: str) -> dict: @cachedmethod(operator.attrgetter("cache"))
@staticmethod
def construct_eval_env(return_type_str: str, imports) -> dict:
""" """
Constructs an evaluation environment with the necessary imports for the return type, Constructs an evaluation environment with the necessary imports for the return type,
taking into account module aliases. taking into account module aliases.
""" """
eval_env: dict = {} eval_env: dict = {}
for import_entry in self.data["imports"]: for import_entry in imports:
if isinstance(import_entry, tuple): # from module import name if isinstance(import_entry, tuple): # from module import name
module, name = import_entry module, name = import_entry
if name in return_type_str: if name in return_type_str:
@ -134,7 +135,7 @@ class CodeParser:
return_type = None return_type = None
if node.returns: if node.returns:
return_type_str = ast.unparse(node.returns) return_type_str = ast.unparse(node.returns)
eval_env = self.construct_eval_env(return_type_str) eval_env = self.construct_eval_env(return_type_str, self.data["imports"])
try: try:
return_type = eval(return_type_str, eval_env) return_type = eval(return_type_str, eval_env)