From 68c6ac8ea24613ef5d9246c7012a0cc9d0e364fe Mon Sep 17 00:00:00 2001 From: gustavoschaedler Date: Wed, 5 Jul 2023 17:23:53 +0100 Subject: [PATCH] Refactor base class removal logic and add dict representation - Extracted the logic for removing unwanted base classes to a separate method called `process_base_classes`. - Created a list `CLASSES_TO_REMOVE` which contains the names of the base classes to be removed. - Updated `to_dict` method to call `process_base_classes` before generating the dict representation of the frontend node. - Added a new method `process_base_classes` which removes the unwanted base classes from the list of base classes. - The method `process_base_classes` updates the `base_classes` list, removing any base classes present in `CLASSES_TO_REMOVE`. This commit enhances the code by separating the concerns of base class removal logic and dict representation generation. --- src/backend/langflow/template/frontend_node/base.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/src/backend/langflow/template/frontend_node/base.py b/src/backend/langflow/template/frontend_node/base.py index cb06789a3..938a4d38e 100644 --- a/src/backend/langflow/template/frontend_node/base.py +++ b/src/backend/langflow/template/frontend_node/base.py @@ -8,6 +8,8 @@ from langflow.template.field.base import TemplateField from langflow.template.template.base import Template from langflow.utils import constants +CLASSES_TO_REMOVE = ["Serializable", "BaseModel"] + class FrontendNode(BaseModel): template: Template @@ -17,7 +19,17 @@ class FrontendNode(BaseModel): display_name: str = "" custom_fields: List[str] = [] + def process_base_classes(self) -> None: + """Removes unwanted base classes from the list of base classes.""" + self.base_classes = [ + base_class + for base_class in self.base_classes + if base_class not in CLASSES_TO_REMOVE + ] + def to_dict(self) -> dict: + """Returns a dict representation of the frontend node.""" + self.process_base_classes() return { self.name: { "template": self.template.to_dict(self.format_field),