Merge pull request #4 from hatchet-dev/feat-finally-fix-mypy

Feat: Finally fixing mypy all the way

Author: mrkaye97

hatchet_sdk/clients/dispatcher/action_listener.py

@@ -2,10 +2,11 @@
 import json
 import time
 from dataclasses import dataclass, field
-from typing import Any, AsyncGenerator, List, Optional
+from typing import Any, AsyncGenerator, AsyncIterable, AsyncIterator, Optional, cast
 
 import grpc
-from grpc._cython import cygrpc
+import grpc.aio
+from grpc._cython import cygrpc  # type: ignore[attr-defined]
 
 from hatchet_sdk.clients.event_ts import Event_ts, read_with_interrupt
 from hatchet_sdk.clients.run_event_listener import (
@@ -40,14 +41,14 @@
 @dataclass
 class GetActionListenerRequest:
     worker_name: str
-    services: List[str]
-    actions: List[str]
+    services: list[str]
+    actions: list[str]
     max_runs: Optional[int] = None
     _labels: dict[str, str | int] = field(default_factory=dict)
 
     labels: dict[str, WorkerLabels] = field(init=False)
 
-    def __post_init__(self):
+    def __post_init__(self) -> None:
         self.labels = {}
 
         for key, value in self._labels.items():
@@ -69,16 +70,16 @@ class Action:
     step_id: str
     step_run_id: str
     action_id: str
-    action_payload: str
     action_type: ActionType
     retry_count: int
+    action_payload: JSONSerializableDict = field(default_factory=dict)
     additional_metadata: JSONSerializableDict = field(default_factory=dict)
 
     child_workflow_index: int | None = None
     child_workflow_key: str | None = None
     parent_workflow_run_id: str | None = None
 
-    def __post_init__(self):
+    def __post_init__(self) -> None:
         if isinstance(self.additional_metadata, str) and self.additional_metadata != "":
             try:
                 self.additional_metadata = json.loads(self.additional_metadata)
@@ -114,11 +115,6 @@ def otel_attributes(self) -> dict[str, Any]:
         )
 
 
-START_STEP_RUN = 0
-CANCEL_STEP_RUN = 1
-START_GET_GROUP_KEY = 2
-
-
 @dataclass
 class ActionListener:
     config: ClientConfig
@@ -131,22 +127,22 @@ class ActionListener:
     last_connection_attempt: float = field(default=0, init=False)
     last_heartbeat_succeeded: bool = field(default=True, init=False)
     time_last_hb_succeeded: float = field(default=9999999999999, init=False)
-    heartbeat_task: Optional[asyncio.Task] = field(default=None, init=False)
+    heartbeat_task: Optional[asyncio.Task[None]] = field(default=None, init=False)
     run_heartbeat: bool = field(default=True, init=False)
     listen_strategy: str = field(default="v2", init=False)
     stop_signal: bool = field(default=False, init=False)
 
     missed_heartbeats: int = field(default=0, init=False)
 
-    def __post_init__(self):
-        self.client = DispatcherStub(new_conn(self.config, False))
-        self.aio_client = DispatcherStub(new_conn(self.config, True))
+    def __post_init__(self) -> None:
+        self.client = DispatcherStub(new_conn(self.config, False))  # type: ignore[no-untyped-call]
+        self.aio_client = DispatcherStub(new_conn(self.config, True))  # type: ignore[no-untyped-call]
         self.token = self.config.token
 
-    def is_healthy(self):
+    def is_healthy(self) -> bool:
         return self.last_heartbeat_succeeded
 
-    async def heartbeat(self):
+    async def heartbeat(self) -> None:
         # send a heartbeat every 4 seconds
         heartbeat_delay = 4
 
@@ -206,7 +202,7 @@ async def heartbeat(self):
                     break
             await asyncio.sleep(heartbeat_delay)
 
-    async def start_heartbeater(self):
+    async def start_heartbeater(self) -> None:
         if self.heartbeat_task is not None:
             return
 
@@ -220,10 +216,10 @@ async def start_heartbeater(self):
                 raise e
         self.heartbeat_task = loop.create_task(self.heartbeat())
 
-    def __aiter__(self):
+    def __aiter__(self) -> AsyncGenerator[Action | None, None]:
         return self._generator()
 
-    async def _generator(self) -> AsyncGenerator[Action, None]:
+    async def _generator(self) -> AsyncGenerator[Action | None, None]:
         listener = None
 
         while not self.stop_signal:
@@ -239,6 +235,10 @@ async def _generator(self) -> AsyncGenerator[Action, None]:
             try:
                 while not self.stop_signal:
                     self.interrupt = Event_ts()
+
+                    if listener is None:
+                        continue
+
                     t = asyncio.create_task(
                         read_with_interrupt(listener, self.interrupt)
                     )
@@ -251,7 +251,10 @@ async def _generator(self) -> AsyncGenerator[Action, None]:
                         )
 
                         t.cancel()
-                        listener.cancel()
+
+                        if listener:
+                            listener.cancel()
+
                         break
 
                     assigned_action = t.result()
@@ -261,20 +264,23 @@ async def _generator(self) -> AsyncGenerator[Action, None]:
                         break
 
                     self.retries = 0
-                    assigned_action: AssignedAction
 
                     # Process the received action
-                    action_type = self.map_action_type(assigned_action.actionType)
+                    action_type = assigned_action.actionType
 
-                    if (
-                        assigned_action.actionPayload is None
-                        or assigned_action.actionPayload == ""
-                    ):
-                        action_payload = None
-                    else:
-                        action_payload = self.parse_action_payload(
-                            assigned_action.actionPayload
+                    action_payload = (
+                        {}
+                        if not assigned_action.actionPayload
+                        else self.parse_action_payload(assigned_action.actionPayload)
+                    )
+
+                    try:
+                        additional_metadata = cast(
+                            dict[str, Any],
+                            json.loads(assigned_action.additional_metadata),
                         )
+                    except json.JSONDecodeError:
+                        additional_metadata = {}
 
                     action = Action(
                         tenant_id=assigned_action.tenantId,
@@ -290,7 +296,7 @@ async def _generator(self) -> AsyncGenerator[Action, None]:
                         action_payload=action_payload,
                         action_type=action_type,
                         retry_count=assigned_action.retryCount,
-                        additional_metadata=assigned_action.additional_metadata,
+                        additional_metadata=additional_metadata,
                         child_workflow_index=assigned_action.child_workflow_index,
                         child_workflow_key=assigned_action.child_workflow_key,
                         parent_workflow_run_id=assigned_action.parent_workflow_run_id,
@@ -324,25 +330,15 @@ async def _generator(self) -> AsyncGenerator[Action, None]:
 
                     self.retries = self.retries + 1
 
-    def parse_action_payload(self, payload: str):
+    def parse_action_payload(self, payload: str) -> JSONSerializableDict:
         try:
-            payload_data = json.loads(payload)
+            return cast(JSONSerializableDict, json.loads(payload))
         except json.JSONDecodeError as e:
             raise ValueError(f"Error decoding payload: {e}")
-        return payload_data
-
-    def map_action_type(self, action_type):
-        if action_type == ActionType.START_STEP_RUN:
-            return START_STEP_RUN
-        elif action_type == ActionType.CANCEL_STEP_RUN:
-            return CANCEL_STEP_RUN
-        elif action_type == ActionType.START_GET_GROUP_KEY:
-            return START_GET_GROUP_KEY
-        else:
-            # logger.error(f"Unknown action type: {action_type}")
-            return None
 
-    async def get_listen_client(self):
+    async def get_listen_client(
+        self,
+    ) -> grpc.aio.UnaryStreamCall[WorkerListenRequest, AssignedAction]:
         current_time = int(time.time())
 
         if (
@@ -370,7 +366,7 @@ async def get_listen_client(self):
                 f"action listener connection interrupted, retrying... ({self.retries}/{DEFAULT_ACTION_LISTENER_RETRY_COUNT})"
             )
 
-        self.aio_client = DispatcherStub(new_conn(self.config, True))
+        self.aio_client = DispatcherStub(new_conn(self.config, True))  # type: ignore[no-untyped-call]
 
         if self.listen_strategy == "v2":
             # we should await for the listener to be established before
@@ -391,11 +387,14 @@ async def get_listen_client(self):
 
         self.last_connection_attempt = current_time
 
-        return listener
+        return cast(
+            grpc.aio.UnaryStreamCall[WorkerListenRequest, AssignedAction], listener
+        )
 
     def cleanup(self) -> None:
         self.run_heartbeat = False
-        self.heartbeat_task.cancel()
+        if self.heartbeat_task is not None:
+            self.heartbeat_task.cancel()
 
         try:
             self.unregister()
@@ -405,9 +404,11 @@ def cleanup(self) -> None:
         if self.interrupt:
             self.interrupt.set()
 
-    def unregister(self):
+    def unregister(self) -> WorkerUnsubscribeRequest:
         self.run_heartbeat = False
-        self.heartbeat_task.cancel()
+
+        if self.heartbeat_task is not None:
+            self.heartbeat_task.cancel()
 
         try:
             req = self.aio_client.Unsubscribe(
@@ -417,6 +418,6 @@ def unregister(self):
             )
             if self.interrupt is not None:
                 self.interrupt.set()
-            return req
+            return cast(WorkerUnsubscribeRequest, req)
         except grpc.RpcError as e:
             raise Exception(f"Failed to unsubscribe: {e}")

hatchet_sdk/clients/dispatcher/dispatcher.py

@@ -1,5 +1,6 @@
 from typing import Any, cast
 
+import grpc.aio
 from google.protobuf.timestamp_pb2 import Timestamp
 
 from hatchet_sdk.clients.dispatcher.action_listener import (
@@ -69,7 +70,7 @@ async def get_action_listener(
 
     async def send_step_action_event(
         self, action: Action, event_type: StepActionEventType, payload: str
-    ) -> Any:
+    ) -> grpc.aio.UnaryUnaryCall[StepActionEvent, ActionEventResponse] | None:
         try:
             return await self._try_send_step_action_event(action, event_type, payload)
         except Exception as e:
@@ -84,12 +85,12 @@ async def send_step_action_event(
                     "Failed to send finished event: " + str(e),
                 )
 
-            return
+            return None
 
     @tenacity_retry
     async def _try_send_step_action_event(
         self, action: Action, event_type: StepActionEventType, payload: str
-    ) -> Any:
+    ) -> grpc.aio.UnaryUnaryCall[StepActionEvent, ActionEventResponse]:
         eventTimestamp = Timestamp()
         eventTimestamp.GetCurrentTime()
 
@@ -105,15 +106,17 @@ async def _try_send_step_action_event(
             eventPayload=payload,
         )
 
-        ## TODO: What does this return?
-        return await self.aio_client.SendStepActionEvent(
-            event,
-            metadata=get_metadata(self.token),
+        return cast(
+            grpc.aio.UnaryUnaryCall[StepActionEvent, ActionEventResponse],
+            await self.aio_client.SendStepActionEvent(
+                event,
+                metadata=get_metadata(self.token),
+            ),
         )
 
     async def send_group_key_action_event(
         self, action: Action, event_type: GroupKeyActionEventType, payload: str
-    ) -> Any:
+    ) -> grpc.aio.UnaryUnaryCall[GroupKeyActionEvent, ActionEventResponse]:
         eventTimestamp = Timestamp()
         eventTimestamp.GetCurrentTime()
 
@@ -128,9 +131,12 @@ async def send_group_key_action_event(
         )
 
         ## TODO: What does this return?
-        return await self.aio_client.SendGroupKeyActionEvent(
-            event,
-            metadata=get_metadata(self.token),
+        return cast(
+            grpc.aio.UnaryUnaryCall[GroupKeyActionEvent, ActionEventResponse],
+            await self.aio_client.SendGroupKeyActionEvent(
+                event,
+                metadata=get_metadata(self.token),
+            ),
         )
 
     def put_overrides_data(self, data: OverridesData) -> ActionEventResponse:

hatchet_sdk/clients/event_ts.py

@@ -1,5 +1,8 @@
 import asyncio
-from typing import Any
+from typing import Any, TypeVar, cast
+
+import grpc.aio
+from grpc._cython import cygrpc  # type: ignore[attr-defined]
 
 
 class Event_ts(asyncio.Event):
@@ -20,9 +23,19 @@ def clear(self) -> None:
         self._loop.call_soon_threadsafe(super().clear)
 
 
-async def read_with_interrupt(listener: Any, interrupt: Event_ts) -> Any:
+TRequest = TypeVar("TRequest")
+TResponse = TypeVar("TResponse")
+
+
+async def read_with_interrupt(
+    listener: grpc.aio.UnaryStreamCall[TRequest, TResponse], interrupt: Event_ts
+) -> TResponse:
     try:
         result = await listener.read()
-        return result
+
+        if result is cygrpc.EOF:
+            raise ValueError("Unexpected EOF")
+
+        return cast(TResponse, result)
     finally:
         interrupt.set()

hatchet_sdk/clients/rest/api_client.py

@@ -97,7 +97,7 @@ async def __aenter__(self):
     async def __aexit__(self, exc_type, exc_value, traceback):
         await self.close()
 
-    async def close(self):
+    async def close(self) -> None:
         await self.rest_client.close()
 
     @property

hatchet_sdk/clients/rest/tenacity_utils.py

@@ -27,7 +27,7 @@ def tenacity_alert_retry(retry_state: tenacity.RetryCallState) -> None:
     )
 
 
-def tenacity_should_retry(ex: Exception) -> bool:
+def tenacity_should_retry(ex: BaseException) -> bool:
     if isinstance(ex, (grpc.aio.AioRpcError, grpc.RpcError)):
         if ex.code() in [
             grpc.StatusCode.UNIMPLEMENTED,