refactor: move type casting from validation function to setter function

VertaAI · Aug 29, 2023 · 93df142 · 93df142 · github-actions · Aug 29, 2023
1 parent 86b1032
commit 93df142
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/client/verta/verta/pipeline/_pipeline_graph.py b/client/verta/verta/pipeline/_pipeline_graph.py
@@ -65,7 +65,7 @@ def set_steps(
         TypeError
             If ``steps`` is not a set of PipelineStep objects.
         """
-        self._steps = self._validate_steps(steps)
+        self._steps = set(self._validate_steps(steps))
         return self.steps
 
     def _validate_steps(
@@ -98,7 +98,7 @@ def _validate_steps(
                     f"individual steps of a PipelineGraph must be type"
                     f" PipelineStep, not {type(step)}."
                 )
-        return set(steps)
+        return steps
 
     @classmethod
     def _from_definition(

diff --git a/client/verta/verta/pipeline/_pipeline_step.py b/client/verta/verta/pipeline/_pipeline_step.py
@@ -167,7 +167,7 @@ def set_predecessors(
             If the provided value for ``steps`` is not a set of PipelineStep objects.
         """
         if steps:
-            self._predecessors = self._validate_predecessors(steps)
+            self._predecessors = set(self._validate_predecessors(steps))
             return self.predecessors
         self._predecessors = set()
         return self.predecessors
@@ -196,7 +196,7 @@ def _validate_predecessors(
                     f"individual predecessors of a PipelineStep must be type"
                     f" PipelineStep, not {type(step)}."
                 )
-        return set(predecessors)
+        return predecessors
 
     def _get_registered_model(self) -> RegisteredModel:
         """Fetch the registered model associated with this step's model version.