Mypy and docstrings

Author: schustmi

src/zenml/execution/pipeline/dynamic/outputs.py

@@ -34,28 +34,61 @@ class OutputArtifact(ArtifactVersionResponse):
 StepRunOutputs = Union[None, OutputArtifact, Tuple[OutputArtifact, ...]]
 
 
-class ArtifactFuture:
-    """Future for a step run output artifact."""
+class _BaseStepRunFuture:
+    """Base step run future."""
 
     def __init__(
-        self, wrapped: Future[StepRunOutputs], invocation_id: str, index: int
+        self,
+        wrapped: Future[StepRunOutputs],
+        invocation_id: str,
+        **kwargs: Any,
     ) -> None:
-        """Initialize the future.
+        """Initialize the dynamic step run future.
 
         Args:
             wrapped: The wrapped future object.
             invocation_id: The invocation ID of the step run.
+            **kwargs: Additional keyword arguments.
         """
         self._wrapped = wrapped
         self._invocation_id = invocation_id
-        self._index = index
+
+    @property
+    def invocation_id(self) -> str:
+        """The step run invocation ID.
+
+        Returns:
+            The step run invocation ID.
+        """
+        return self._invocation_id
 
     def _wait(self) -> None:
+        """Wait for the step run future to complete."""
         self._wrapped.result()
 
+
+class ArtifactFuture(_BaseStepRunFuture):
+    """Future for a step run output artifact."""
+
+    def __init__(
+        self, wrapped: Future[StepRunOutputs], invocation_id: str, index: int
+    ) -> None:
+        """Initialize the future.
+
+        Args:
+            wrapped: The wrapped future object.
+            invocation_id: The invocation ID of the step run.
+            index: The index of the output artifact.
+        """
+        super().__init__(wrapped=wrapped, invocation_id=invocation_id)
+        self._index = index
+
     def result(self) -> OutputArtifact:
         """Get the step run output artifact.
 
+        Raises:
+            RuntimeError: If the future returned an invalid output.
+
         Returns:
             The step run output artifact.
         """
@@ -66,7 +99,8 @@ def result(self) -> OutputArtifact:
             return result[self._index]
         else:
             raise RuntimeError(
-                f"Step {self._invocation_id} returned an invalid output: {result}"
+                f"Step {self._invocation_id} returned an invalid output: "
+                f"{result}."
             )
 
     def load(self) -> Any:
@@ -78,7 +112,7 @@ def load(self) -> Any:
         return self.result().load()
 
 
-class StepRunOutputsFuture:
+class StepRunOutputsFuture(_BaseStepRunFuture):
     """Future for a step run output."""
 
     def __init__(
@@ -92,14 +126,11 @@ def __init__(
         Args:
             wrapped: The wrapped future object.
             invocation_id: The invocation ID of the step run.
+            output_keys: The output keys of the step run.
         """
-        self._wrapped = wrapped
-        self._invocation_id = invocation_id
+        super().__init__(wrapped=wrapped, invocation_id=invocation_id)
         self._output_keys = output_keys
 
-    def _wait(self) -> None:
-        self._wrapped.result()
-
     def artifacts(self) -> StepRunOutputs:
         """Get the step run output artifacts.
 

src/zenml/execution/pipeline/dynamic/runner.py

@@ -41,6 +41,7 @@
     StepRunFuture,
     StepRunOutputs,
     StepRunOutputsFuture,
+    _BaseStepRunFuture,
 )
 from zenml.execution.pipeline.dynamic.run_context import (
     DynamicPipelineRunContext,
@@ -292,22 +293,22 @@ def compile_dynamic_step_invocation(
         pipeline: The dynamic pipeline.
         step: The step to compile.
         id: Custom invocation ID.
-        upstream_steps: The upstream steps.
-        inputs: Inputs to the step function.
-        default_parameters: Default parameters of the step function.
+        args: The arguments for the step function.
+        kwargs: The keyword arguments for the step function.
+        after: The step run output futures to wait for.
 
     Returns:
         The compiled step.
     """
     upstream_steps = set()
 
-    if isinstance(after, StepRunFuture):
+    if isinstance(after, _BaseStepRunFuture):
         after._wait()
-        upstream_steps.add(after._invocation_id)
+        upstream_steps.add(after.invocation_id)
     elif isinstance(after, Sequence):
         for item in after:
             item._wait()
-            upstream_steps.add(item._invocation_id)
+            upstream_steps.add(item.invocation_id)
 
     def _await_and_validate_input(input: Any) -> Any:
         if isinstance(input, StepRunOutputsFuture):

src/zenml/execution/pipeline/utils.py

@@ -62,7 +62,11 @@ def should_prevent_pipeline_execution() -> bool:
 
 @contextmanager
 def prevent_pipeline_execution() -> Generator[None, None, None]:
-    """Context manager to prevent pipeline execution."""
+    """Context manager to prevent pipeline execution.
+
+    Yields:
+        None.
+    """
     with env_utils.temporary_environment(
         {ENV_ZENML_PREVENT_PIPELINE_EXECUTION: "True"}
     ):

src/zenml/execution/step/utils.py

@@ -48,6 +48,10 @@ def launch_step(
         orchestrator_run_id: The orchestrator run ID.
         retry: Whether to retry the step if it fails.
 
+    Raises:
+        RunStoppedException: If the run was stopped.
+        BaseException: If the step failed all retries.
+
     Returns:
         The step run response.
     """

src/zenml/orchestrators/base_orchestrator.py

@@ -433,10 +433,6 @@ def run_step(
 
         Args:
             step: The step to run.
-
-        Raises:
-            RunStoppedException: If the run was stopped.
-            BaseException: If the step failed all retries.
         """
         from zenml.execution.step.utils import launch_step
 
@@ -482,6 +478,10 @@ def launch_dynamic_step(
             step_run_info: The step run information.
             environment: The environment variables to set in the execution
                 environment.
+
+        Raises:
+            NotImplementedError: If the orchestrator does not implement this
+                method.
         """
         raise NotImplementedError(
             "Launching dynamic steps is not implemented for "

src/zenml/orchestrators/local/local_orchestrator.py

@@ -193,7 +193,18 @@ def submit_dynamic_pipeline(
         environment: Dict[str, str],
         placeholder_run: Optional["PipelineRunResponse"] = None,
     ) -> Optional[SubmissionResult]:
-        """Submits a dynamic pipeline to the orchestrator."""
+        """Submits a dynamic pipeline to the orchestrator.
+
+        Args:
+            snapshot: The pipeline snapshot to submit.
+            stack: The stack the pipeline will run on.
+            environment: Environment variables to set in the orchestration
+                environment.
+            placeholder_run: An optional placeholder run.
+
+        Returns:
+            Optional submission result.
+        """
         from zenml.execution.pipeline.dynamic.runner import (
             DynamicPipelineRunner,
         )

src/zenml/orchestrators/publish_utils.py

@@ -125,6 +125,9 @@ def publish_successful_pipeline_run(
 
     Args:
         pipeline_run_id: The ID of the pipeline run to update.
+
+    Returns:
+        The updated pipeline run.
     """
     return Client().zen_store.update_run(
         run_id=pipeline_run_id,

src/zenml/orchestrators/step_launcher.py

@@ -252,6 +252,9 @@ def launch(self) -> StepRunResponse:
         Raises:
             RunStoppedException: If the pipeline run is stopped by the user.
             BaseException: If the step preparation or execution fails.
+
+        Returns:
+            The step run response.
         """
         publish_utils.step_exception_info.set(None)
         pipeline_run, run_was_created = self._create_or_reuse_run()

src/zenml/pipelines/dynamic/pipeline_definition.py

@@ -103,7 +103,11 @@ def is_dynamic(self) -> bool:
         return True
 
     def _prepare_invocations(self, **kwargs: Any) -> None:
-        """Prepares the invocations of the pipeline."""
+        """Prepares the invocations of the pipeline.
+
+        Args:
+            **kwargs: Keyword arguments.
+        """
         for step in self._depends_on:
             self.add_step_invocation(
                 step,

src/zenml/pipelines/pipeline_decorator.py

@@ -106,6 +106,7 @@ def pipeline(
         _func: The decorated function.
         name: The name of the pipeline. If left empty, the name of the
             decorated function will be used as a fallback.
+        depends_on: The steps that this pipeline depends on.
         enable_cache: Whether to use caching or not.
         enable_artifact_metadata: Whether to enable artifact metadata or not.
         enable_step_logs: If step logs should be enabled for this pipeline.