Task:4019 Upgrade type hints to 3.10+ style

scripts/pyupgrade.sh

@@ -0,0 +1,18 @@
+#!/bin/bash
+
+# install helper packages
+
+uv pip install pyupgrade
+uv pip install autoflake
+
+# run py-upgrade recursively to all nested .py files
+
+find src -name '*.py' -print0 | xargs -0 -n1 python -m pyupgrade --py310-plus
+
+# run autoflake to remove dangling imports
+
+autoflake --remove-all-unused-imports --in-place --recursive src
+
+# run ruff check to verify imports are fixed
+
+ruff check core --select F401

src/zenml/actions/base_action.py

@@ -14,7 +14,7 @@
 """Base implementation of actions."""
 
 from abc import ABC, abstractmethod
-from typing import Any, ClassVar, Dict, Optional, Type
+from typing import Any, ClassVar
 
 from zenml.enums import PluginType
 from zenml.event_hub.base_event_hub import BaseEventHub
@@ -56,10 +56,10 @@ class BaseActionFlavor(BasePluginFlavor, ABC):
     TYPE: ClassVar[PluginType] = PluginType.ACTION
 
     # Action specific
-    ACTION_CONFIG_CLASS: ClassVar[Type[ActionConfig]]
+    ACTION_CONFIG_CLASS: ClassVar[type[ActionConfig]]
 
     @classmethod
-    def get_action_config_schema(cls) -> Dict[str, Any]:
+    def get_action_config_schema(cls) -> dict[str, Any]:
         """The config schema for a flavor.
 
         Returns:
@@ -97,9 +97,9 @@ def get_flavor_response_model(cls, hydrate: bool) -> ActionFlavorResponse:
 class BaseActionHandler(BasePlugin, ABC):
     """Implementation for an action handler."""
 
-    _event_hub: Optional[BaseEventHub] = None
+    _event_hub: BaseEventHub | None = None
 
-    def __init__(self, event_hub: Optional[BaseEventHub] = None) -> None:
+    def __init__(self, event_hub: BaseEventHub | None = None) -> None:
         """Event source handler initialization.
 
         Args:
@@ -122,7 +122,7 @@ def __init__(self, event_hub: Optional[BaseEventHub] = None) -> None:
 
     @property
     @abstractmethod
-    def config_class(self) -> Type[ActionConfig]:
+    def config_class(self) -> type[ActionConfig]:
         """Returns the `BasePluginConfig` config.
 
         Returns:
@@ -131,7 +131,7 @@ def config_class(self) -> Type[ActionConfig]:
 
     @property
     @abstractmethod
-    def flavor_class(self) -> Type[BaseActionFlavor]:
+    def flavor_class(self) -> type[BaseActionFlavor]:
         """Returns the flavor class of the plugin.
 
         Returns:
@@ -173,7 +173,7 @@ def set_event_hub(self, event_hub: BaseEventHub) -> None:
 
     def event_hub_callback(
         self,
-        config: Dict[str, Any],
+        config: dict[str, Any],
         trigger_execution: TriggerExecutionResponse,
         auth_context: AuthContext,
     ) -> None:
@@ -442,7 +442,7 @@ def get_action(
         return action
 
     def validate_action_configuration(
-        self, action_config: Dict[str, Any]
+        self, action_config: dict[str, Any]
     ) -> ActionConfig:
         """Validate and return the action configuration.
 
@@ -464,7 +464,7 @@ def extract_resources(
         self,
         action_config: ActionConfig,
         hydrate: bool = False,
-    ) -> Dict[ResourceType, BaseResponse[Any, Any, Any]]:
+    ) -> dict[ResourceType, BaseResponse[Any, Any, Any]]:
         """Extract related resources for this action.
 
         Args:
@@ -504,7 +504,6 @@ def _validate_action_request(
             action: Action request.
             config: Action configuration instantiated from the request.
         """
-        pass
 
     def _process_action_request(
         self, action: ActionResponse, config: ActionConfig
@@ -529,7 +528,6 @@ def _process_action_request(
             action: Newly created action
             config: Action configuration instantiated from the response.
         """
-        pass
 
     def _validate_action_update(
         self,
@@ -566,7 +564,6 @@ def _validate_action_update(
             config_update: Action configuration instantiated from the
                 updated action.
         """
-        pass
 
     def _process_action_update(
         self,
@@ -599,13 +596,12 @@ def _process_action_update(
             previous_config: Action configuration instantiated from the
                 original action.
         """
-        pass
 
     def _process_action_delete(
         self,
         action: ActionResponse,
         config: ActionConfig,
-        force: Optional[bool] = False,
+        force: bool | None = False,
     ) -> None:
         """Process an action before it is deleted from the database.
 
@@ -625,7 +621,6 @@ def _process_action_delete(
             config: Action configuration before the deletion.
             force: Whether to force deprovision the action.
         """
-        pass
 
     def _process_action_response(
         self, action: ActionResponse, config: ActionConfig
@@ -650,7 +645,6 @@ def _process_action_response(
             action: Action response.
             config: Action configuration instantiated from the response.
         """
-        pass
 
     def _populate_action_response_resources(
         self,

src/zenml/actions/pipeline_run/pipeline_run_action.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """Example file of what an action Plugin could look like."""
 
-from typing import Any, ClassVar, Dict, Optional, Type
+from typing import Any, ClassVar
 from uuid import UUID
 
 from zenml.actions.base_action import (
@@ -47,7 +47,7 @@ class PipelineRunActionConfiguration(ActionConfig):
     """Configuration class to configure a pipeline run action."""
 
     snapshot_id: UUID
-    run_config: Optional[PipelineRunConfiguration] = None
+    run_config: PipelineRunConfiguration | None = None
 
 
 # -------------------- Pipeline Run Plugin -----------------------------------
@@ -57,7 +57,7 @@ class PipelineRunActionHandler(BaseActionHandler):
     """Action handler for running pipelines."""
 
     @property
-    def config_class(self) -> Type[PipelineRunActionConfiguration]:
+    def config_class(self) -> type[PipelineRunActionConfiguration]:
         """Returns the `BasePluginConfig` config.
 
         Returns:
@@ -66,7 +66,7 @@ def config_class(self) -> Type[PipelineRunActionConfiguration]:
         return PipelineRunActionConfiguration
 
     @property
-    def flavor_class(self) -> Type[BaseActionFlavor]:
+    def flavor_class(self) -> type[BaseActionFlavor]:
         """Returns the flavor class of the plugin.
 
         Returns:
@@ -169,7 +169,7 @@ def extract_resources(
         self,
         action_config: ActionConfig,
         hydrate: bool = False,
-    ) -> Dict[ResourceType, BaseResponse[Any, Any, Any]]:
+    ) -> dict[ResourceType, BaseResponse[Any, Any, Any]]:
         """Extract related resources for this action.
 
         Args:
@@ -196,7 +196,7 @@ def extract_resources(
                 f"No snapshot found with id {action_config.snapshot_id}."
             )
 
-        resources: Dict[ResourceType, BaseResponse[Any, Any, Any]] = {
+        resources: dict[ResourceType, BaseResponse[Any, Any, Any]] = {
             ResourceType.PIPELINE_SNAPSHOT: snapshot
         }
 
@@ -217,11 +217,11 @@ class PipelineRunActionFlavor(BaseActionFlavor):
 
     FLAVOR: ClassVar[str] = "builtin"
     SUBTYPE: ClassVar[PluginSubType] = PluginSubType.PIPELINE_RUN
-    PLUGIN_CLASS: ClassVar[Type[PipelineRunActionHandler]] = (
+    PLUGIN_CLASS: ClassVar[type[PipelineRunActionHandler]] = (
         PipelineRunActionHandler
     )
 
     # EventPlugin specific
-    ACTION_CONFIG_CLASS: ClassVar[Type[PipelineRunActionConfiguration]] = (
+    ACTION_CONFIG_CLASS: ClassVar[type[PipelineRunActionConfiguration]] = (
         PipelineRunActionConfiguration
     )

src/zenml/alerter/base_alerter.py

@@ -14,7 +14,7 @@
 """Base class for all ZenML alerters."""
 
 from abc import ABC
-from typing import Optional, Type, cast
+from typing import cast
 
 from pydantic import BaseModel
 
@@ -44,7 +44,7 @@ def config(self) -> BaseAlerterConfig:
         return cast(BaseAlerterConfig, self._config)
 
     def post(
-        self, message: str, params: Optional[BaseAlerterStepParameters] = None
+        self, message: str, params: BaseAlerterStepParameters | None = None
     ) -> bool:
         """Post a message to a chat service.
 
@@ -58,7 +58,7 @@ def post(
         return True
 
     def ask(
-        self, question: str, params: Optional[BaseAlerterStepParameters] = None
+        self, question: str, params: BaseAlerterStepParameters | None = None
     ) -> bool:
         """Post a message to a chat service and wait for approval.
 
@@ -88,7 +88,7 @@ def type(self) -> StackComponentType:
         return StackComponentType.ALERTER
 
     @property
-    def config_class(self) -> Type[BaseAlerterConfig]:
+    def config_class(self) -> type[BaseAlerterConfig]:
         """Returns BaseAlerterConfig class.
 
         Returns:
@@ -97,7 +97,7 @@ def config_class(self) -> Type[BaseAlerterConfig]:
         return BaseAlerterConfig
 
     @property
-    def implementation_class(self) -> Type[BaseAlerter]:
+    def implementation_class(self) -> type[BaseAlerter]:
         """Implementation class.
 
         Returns:

src/zenml/analytics/__init__.py

@@ -13,7 +13,7 @@
 #  permissions and limitations under the License.
 """The 'analytics' module of ZenML."""
 from contextvars import ContextVar
-from typing import Any, Dict, Optional, TYPE_CHECKING
+from typing import Any, TYPE_CHECKING
 from uuid import UUID
 
 from zenml.enums import SourceContextTypes
@@ -27,7 +27,7 @@
 
 
 def identify(  # type: ignore[return]
-    metadata: Optional[Dict[str, Any]] = None
+    metadata: dict[str, Any] | None = None
 ) -> bool:
     """Attach metadata to user directly.
 
@@ -64,7 +64,7 @@ def alias(user_id: UUID, previous_id: UUID) -> bool:  # type: ignore[return]
 
 def group(  # type: ignore[return]
     group_id: UUID,
-    group_metadata: Optional[Dict[str, Any]] = None,
+    group_metadata: dict[str, Any] | None = None,
 ) -> bool:
     """Attach metadata to a segment group.
 
@@ -83,7 +83,7 @@ def group(  # type: ignore[return]
 
 def track(  # type: ignore[return]
     event: "AnalyticsEvent",
-    metadata: Optional[Dict[str, Any]] = None,
+    metadata: dict[str, Any] | None = None,
 ) -> bool:
     """Track segment event if user opted-in.
 

src/zenml/analytics/client.py

@@ -15,7 +15,7 @@
 
 import json
 import logging
-from typing import Any, Dict, Optional, Tuple
+from typing import Any
 from uuid import UUID
 
 from zenml.analytics.enums import AnalyticsEvent
@@ -26,7 +26,7 @@
 logger = logging.getLogger(__name__)
 
 
-class Client(object):
+class Client:
     """The client class for ZenML analytics."""
 
     def __init__(self, send: bool = True, timeout: int = 15) -> None:
@@ -40,8 +40,8 @@ def __init__(self, send: bool = True, timeout: int = 15) -> None:
         self.timeout = timeout
 
     def identify(
-        self, user_id: UUID, traits: Optional[Dict[Any, Any]]
-    ) -> Tuple[bool, str]:
+        self, user_id: UUID, traits: dict[Any, Any] | None
+    ) -> tuple[bool, str]:
         """Method to identify a user with given traits.
 
         Args:
@@ -59,7 +59,7 @@ def identify(
         }
         return self._enqueue(json.dumps(msg, cls=AnalyticsEncoder))
 
-    def alias(self, user_id: UUID, previous_id: UUID) -> Tuple[bool, str]:
+    def alias(self, user_id: UUID, previous_id: UUID) -> tuple[bool, str]:
         """Method to alias user IDs.
 
         Args:
@@ -81,8 +81,8 @@ def track(
         self,
         user_id: UUID,
         event: "AnalyticsEvent",
-        properties: Optional[Dict[Any, Any]],
-    ) -> Tuple[bool, str]:
+        properties: dict[Any, Any] | None,
+    ) -> tuple[bool, str]:
         """Method to track events.
 
         Args:
@@ -103,8 +103,8 @@ def track(
         return self._enqueue(json.dumps(msg, cls=AnalyticsEncoder))
 
     def group(
-        self, user_id: UUID, group_id: UUID, traits: Optional[Dict[Any, Any]]
-    ) -> Tuple[bool, str]:
+        self, user_id: UUID, group_id: UUID, traits: dict[Any, Any] | None
+    ) -> tuple[bool, str]:
         """Method to group users.
 
         Args:
@@ -124,7 +124,7 @@ def group(
         }
         return self._enqueue(json.dumps(msg, cls=AnalyticsEncoder))
 
-    def _enqueue(self, msg: str) -> Tuple[bool, str]:
+    def _enqueue(self, msg: str) -> tuple[bool, str]:
         """Method to queue messages to be sent.
 
         Args: