Linting

Author: schustmi

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

@@ -160,6 +160,18 @@ def load(self) -> Any:
             raise ValueError(f"Invalid step run output: {result}")
 
     def __getitem__(self, key: Union[str, int]) -> ArtifactFuture:
+        """Get an artifact future by key or index.
+
+        Args:
+            key: The key or index of the artifact future.
+
+        Raises:
+            ValueError: If the key or index is of an invalid type.
+            IndexError: If the index is out of range.
+
+        Returns:
+            The artifact future.
+        """
         if isinstance(key, str):
             index = self._output_keys.index(key)
         elif isinstance(key, int):
@@ -177,6 +189,14 @@ def __getitem__(self, key: Union[str, int]) -> ArtifactFuture:
         )
 
     def __iter__(self) -> Any:
+        """Iterate over the artifact futures.
+
+        Raises:
+            ValueError: If the step does not return any outputs.
+
+        Yields:
+            The artifact futures.
+        """
         if not self._output_keys:
             raise ValueError(
                 f"Step {self._invocation_id} does not return any outputs."
@@ -190,6 +210,11 @@ def __iter__(self) -> Any:
             )
 
     def __len__(self) -> int:
+        """Get the number of artifact futures.
+
+        Returns:
+            The number of artifact futures.
+        """
         return len(self._output_keys)
 
 

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

@@ -14,7 +14,9 @@
 """Dynamic pipeline run context."""
 
 import contextvars
-from typing import TYPE_CHECKING, Self
+from typing import TYPE_CHECKING
+
+from typing_extensions import Self
 
 from zenml.utils import context_utils
 

src/zenml/integrations/kubernetes/orchestrators/kubernetes_orchestrator.py

@@ -685,7 +685,7 @@ def _submit_orchestrator_job(
         args: List[str],
         environment: Dict[str, str],
         placeholder_run: Optional["PipelineRunResponse"] = None,
-    ) -> None:
+    ) -> Optional[SubmissionResult]:
         """Submits an orchestrator job to Kubernetes.
 
         Args:
@@ -830,7 +830,7 @@ def launch_dynamic_step(
         )
 
         logger.info(
-            f"Launching job for step `%s`.",
+            "Launching job for step `%s`.",
             step_run_info.pipeline_step_name,
         )
 

src/zenml/logging/step_logging.py

@@ -940,6 +940,20 @@ def setup_pipeline_logging(
     run_id: Optional[UUID] = None,
     logs_response: Optional[LogsResponse] = None,
 ) -> Generator[Optional[LogsRequest], None, None]:
+    """Set up logging for a pipeline run.
+
+    Args:
+        source: The log source.
+        snapshot: The snapshot of the pipeline run.
+        run_id: The ID of the pipeline run.
+        logs_response: The logs response to continue from.
+
+    Raises:
+        Exception: If updating the run with the logs request fails.
+
+    Yields:
+        The logs request.
+    """
     logging_enabled = True
 
     if handle_bool_env_var(ENV_ZENML_DISABLE_PIPELINE_LOGS_STORAGE, False):

src/zenml/pipelines/compilation_context.py

@@ -1,5 +1,22 @@
+#  Copyright (c) ZenML GmbH 2025. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at:
+#
+#       https://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+#  or implied. See the License for the specific language governing
+#  permissions and limitations under the License.
+"""Pipeline compilation context."""
+
 import contextvars
-from typing import TYPE_CHECKING, Self
+from typing import TYPE_CHECKING
+
+from typing_extensions import Self
 
 from zenml.utils import context_utils
 

src/zenml/step_operators/step_operator_entrypoint_configuration.py

@@ -36,6 +36,12 @@ class StepOperatorEntrypointConfiguration(StepEntrypointConfiguration):
     """Base class for step operator entrypoint configurations."""
 
     def __init__(self, *args: Any, **kwargs: Any) -> None:
+        """Initialize the step operator entrypoint configuration.
+
+        Args:
+            *args: The arguments to pass to the superclass.
+            **kwargs: The keyword arguments to pass to the superclass.
+        """
         super().__init__(*args, **kwargs)
         self._step_run: Optional["StepRunResponse"] = None
 

src/zenml/steps/base_step.py

@@ -633,6 +633,22 @@ def submit(
         after: Union["StepRunFuture", Sequence["StepRunFuture"], None] = None,
         **kwargs: Any,
     ) -> "StepRunOutputsFuture":
+        """Submit the step to run concurrently in a separate thread.
+
+        Args:
+            *args: The arguments to pass to the step function.
+            id: The invocation ID of the step.
+            after: The step run output futures to wait for before executing the
+                step.
+            **kwargs: The keyword arguments to pass to the step function.
+
+        Raises:
+            RuntimeError: If this method is called outside of a dynamic
+                pipeline.
+
+        Returns:
+            The step run output future.
+        """
         from zenml.execution.pipeline.dynamic.run_context import (
             DynamicPipelineRunContext,
         )

src/zenml/steps/step_context.py

@@ -21,11 +21,12 @@
     List,
     Mapping,
     Optional,
-    Self,
     Sequence,
     Type,
 )
 
+from typing_extensions import Self
+
 from zenml.exceptions import StepContextError
 from zenml.logger import get_logger
 from zenml.utils import context_utils

src/zenml/utils/context_utils.py

@@ -16,7 +16,9 @@
 import contextvars
 import threading
 from contextvars import ContextVar
-from typing import Any, ClassVar, Generic, List, Optional, Self, TypeVar
+from typing import Any, ClassVar, Generic, List, Optional, TypeVar
+
+from typing_extensions import Self
 
 T = TypeVar("T")