From d16d916952a728a834ef1677efbf82f3fa5391e5 Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Wed, 9 Aug 2023 15:28:39 -0300 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix(custom=5Fcomponent.py):=20im?= =?UTF-8?q?prove=20parsing=20of=20return=20type=20in=20CustomComponent=20c?= =?UTF-8?q?lass?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ✨ feat(custom_component.py): add support for parsing return type when it is a Union of types in CustomComponent class --- .../interface/custom/custom_component.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/src/backend/langflow/interface/custom/custom_component.py b/src/backend/langflow/interface/custom/custom_component.py index 8c0b2537a..15a414050 100644 --- a/src/backend/langflow/interface/custom/custom_component.py +++ b/src/backend/langflow/interface/custom/custom_component.py @@ -118,15 +118,15 @@ class CustomComponent(Component, extra=Extra.allow): build_method = build_methods[0] return_type = build_method["return_type"] - # It could be a type or a Union[type1, type2] - if "Union" in return_type: - return_type = ( - return_type.replace("Union", "").replace("[", "").replace("]", "") - ) - return_type = return_type.split(",") - return_type = [item.strip() for item in return_type] - return [item for item in return_type if item in self.return_type_valid_list] - return build_method["return_type"] + # If the return type is not a Union, then we just return it as a list + if "Union" not in return_type: + return [return_type] if return_type in self.return_type_valid_list else [] + + # If the return type is a Union, then we need to parse it + return_type = return_type.replace("Union", "").replace("[", "").replace("]", "") + return_type = return_type.split(",") + return_type = [item.strip() for item in return_type] + return [item for item in return_type if item in self.return_type_valid_list] @property def get_main_class_name(self):