From 1cd651e30042d2d2bb349805d04d1c03b0089238 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Fri, 29 Mar 2024 22:16:11 -0300 Subject: [PATCH] Refactor code_parser/utils.py and custom_component/component.py --- .../base/langflow/interface/custom/code_parser/utils.py | 5 +++-- .../langflow/interface/custom/custom_component/component.py | 4 ---- .../interface/custom/custom_component/custom_component.py | 5 +++-- 3 files changed, 6 insertions(+), 8 deletions(-) diff --git a/src/backend/base/langflow/interface/custom/code_parser/utils.py b/src/backend/base/langflow/interface/custom/code_parser/utils.py index d9b9def26..d73a109a6 100644 --- a/src/backend/base/langflow/interface/custom/code_parser/utils.py +++ b/src/backend/base/langflow/interface/custom/code_parser/utils.py @@ -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__) diff --git a/src/backend/base/langflow/interface/custom/custom_component/component.py b/src/backend/base/langflow/interface/custom/custom_component/component.py index 20b23980d..470ebcde8 100644 --- a/src/backend/base/langflow/interface/custom/custom_component/component.py +++ b/src/backend/base/langflow/interface/custom/custom_component/component.py @@ -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): diff --git a/src/backend/base/langflow/interface/custom/custom_component/custom_component.py b/src/backend/base/langflow/interface/custom/custom_component/custom_component.py index c61b6ffb0..409693681 100644 --- a/src/backend/base/langflow/interface/custom/custom_component/custom_component.py +++ b/src/backend/base/langflow/interface/custom/custom_component/custom_component.py @@ -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