refactor: Improve handling of template_dicts inputs in Vertex class

Refactor the Vertex class to enhance the handling of template_dicts inputs. Instead of using list comprehensions, explicit loops are now used to append the required and optional inputs. This change improves the readability and maintainability of the code.
This commit is contained in:
ogabrielluiz 2024-06-13 16:51:27 -03:00
commit 96511fef4a

View file

@ -98,10 +98,15 @@ class Component(CustomComponent):
raw = self.status raw = self.status
if hasattr(raw, "data") and raw is not None: if hasattr(raw, "data") and raw is not None:
raw = raw.data raw = raw.data
if raw is None:
raw = custom_repr
elif hasattr(raw, "model_dump") and raw is not None: elif hasattr(raw, "model_dump") and raw is not None:
raw = raw.model_dump() raw = raw.model_dump()
artifact_type = get_artifact_type(self.status, result) if raw is None and isinstance(result, (dict, Data, str)):
raw = result.data if isinstance(result, Data) else result
artifact_type = get_artifact_type(self.repr_value or raw, result)
raw = post_process_raw(raw, artifact_type) raw = post_process_raw(raw, artifact_type)
artifact = {"repr": custom_repr, "raw": raw, "type": artifact_type} artifact = {"repr": custom_repr, "raw": raw, "type": artifact_type}
_artifacts[output.name] = artifact _artifacts[output.name] = artifact
@ -110,21 +115,15 @@ class Component(CustomComponent):
return _results, _artifacts return _results, _artifacts
def custom_repr(self): def custom_repr(self):
# ! Temporary REPR if self.repr_value == "":
# Since all are dict, yaml.dump them self.repr_value = self.status
if isinstance(self._results, dict): if isinstance(self.repr_value, dict):
_build_results = recursive_serialize_or_str(self._results) return yaml.dump(self.repr_value)
try: if isinstance(self.repr_value, str):
custom_repr = yaml.dump(_build_results) return self.repr_value
except Exception as e: if isinstance(self.repr_value, BaseModel) and not isinstance(self.repr_value, Data):
logger.error(f"Error while dumping build_result: {e}") return str(self.repr_value)
custom_repr = str(self._results) return self.repr_value
if custom_repr is None and isinstance(self._results, (dict, Data, str)):
custom_repr = self._results
if not isinstance(custom_repr, str):
custom_repr = str(custom_repr)
return custom_repr
def build_inputs(self, user_id: Optional[Union[str, UUID]] = None): def build_inputs(self, user_id: Optional[Union[str, UUID]] = None):
""" """