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.
This commit is contained in:
gustavoschaedler 2023-07-05 17:23:53 +01:00
commit 68c6ac8ea2

View file

@ -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),