From 17c6ed1db3b167638e8d3ce6d2bd672f42f7503a Mon Sep 17 00:00:00 2001 From: Gabriel Luiz Freitas Almeida Date: Mon, 11 Dec 2023 12:30:03 -0300 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20feat(base.py):=20add=20methods=20to?= =?UTF-8?q?=20get,=20update,=20and=20upsert=20fields=20in=20the=20Template?= =?UTF-8?q?=20class=20for=20easier=20field=20manipulation=20=F0=9F=94=A7?= =?UTF-8?q?=20chore(base.py):=20refactor=20Template=20class=20to=20improve?= =?UTF-8?q?=20code=20readability=20and=20maintainability?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../langflow/template/template/base.py | 25 ++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/src/backend/langflow/template/template/base.py b/src/backend/langflow/template/template/base.py index bb5d35375..d7632e239 100644 --- a/src/backend/langflow/template/template/base.py +++ b/src/backend/langflow/template/template/base.py @@ -1,8 +1,9 @@ from typing import Callable, Union +from pydantic import BaseModel, model_serializer + from langflow.template.field.base import TemplateField from langflow.utils.constants import DIRECT_TYPES -from pydantic import BaseModel, model_serializer class Template(BaseModel): @@ -39,3 +40,25 @@ class Template(BaseModel): def add_field(self, field: TemplateField) -> None: self.fields.append(field) + + def get_field(self, field_name: str) -> TemplateField: + """Returns the field with the given name.""" + field = next((field for field in self.fields if field.name == field_name), None) + if field is None: + raise ValueError(f"Field {field_name} not found in template {self.type_name}") + return field + + def update_field(self, field_name: str, field: TemplateField) -> None: + """Updates the field with the given name.""" + for idx, template_field in enumerate(self.fields): + if template_field.name == field_name: + self.fields[idx] = field + return + raise ValueError(f"Field {field_name} not found in template {self.type_name}") + + def upsert_field(self, field_name: str, field: TemplateField) -> None: + """Updates the field with the given name or adds it if it doesn't exist.""" + try: + self.update_field(field_name, field) + except ValueError: + self.add_field(field)