refactor: Update stop method in CustomComponent to allow for optional output_name parameter

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-06-20 21:51:45 -03:00
commit 5561b4ec52

View file

@ -91,7 +91,7 @@ class CustomComponent(BaseComponent):
except Exception as e:
raise ValueError(f"Error updating state: {e}")
def stop(self, output_name: str | None):
def stop(self, output_name: str | None = None):
if not output_name and self.vertex and len(self.vertex.outputs) == 1:
output_name = self.vertex.outputs[0]["name"]
else: