Refactor code_parser/utils.py and custom_component/component.py

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-03-29 22:16:11 -03:00
commit 1cd651e300
3 changed files with 6 additions and 8 deletions

View file

@ -14,11 +14,10 @@ def extract_inner_type(return_type: str) -> str:
def extract_inner_type_from_generic_alias(return_type: GenericAlias) -> Any:
"""
Extracts the inner type from a type hint that is a list.
Extracts the inner type from a type hint that is a list or a Optional.
"""
if return_type.__origin__ == list:
return list(return_type.__args__)
return return_type
@ -36,4 +35,6 @@ def extract_union_types_from_generic_alias(return_type: GenericAlias) -> list:
"""
Extracts the inner type from a type hint that is a Union.
"""
if isinstance(return_type, list):
return [_inner_arg for _type in return_type for _inner_arg in _type.__args__]
return list(return_type.__args__)

View file

@ -34,10 +34,6 @@ class Component:
if key == "user_id":
setattr(self, "_user_id", value)
else:
if key == "code" and "from langflow import CustomComponent" in value:
value = value.replace(
"from langflow import CustomComponent", "from langflow.custom import CustomComponent"
)
setattr(self, key, value)
def __setattr__(self, key, value):

View file

@ -318,9 +318,10 @@ class CustomComponent(Component):
return_type = extract_inner_type_from_generic_alias(return_type)
# If the return type is not a Union, then we just return it as a list
if not hasattr(return_type, "__origin__") or return_type.__origin__ != Union:
inner_type = return_type[0] if isinstance(return_type, list) else return_type
if not hasattr(inner_type, "__origin__") or inner_type.__origin__ != Union:
return return_type if isinstance(return_type, list) else [return_type]
# If the return type is a Union, then we need to parse itx
# If the return type is a Union, then we need to parse it
return_type = extract_union_types_from_generic_alias(return_type)
return return_type