fix: remove some Any types

Author: mrkaye97

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

@@ -2,6 +2,7 @@
 from typing import Any, TypeVar, cast
 
 import grpc.aio
+from grpc._cython import cygrpc  # type: ignore[attr-defined]
 
 
 class Event_ts(asyncio.Event):
@@ -28,8 +29,13 @@ def clear(self) -> None:
 
 async def read_with_interrupt(
     listener: grpc.aio.UnaryStreamCall[TRequest, TResponse], interrupt: Event_ts
-) -> Any:
+) -> TResponse:
     try:
-        return cast(Any, await listener.read())
+        result = await listener.read()
+
+        if result is cygrpc.EOF:
+            raise ValueError("Unexpected EOF")
+
+        return cast(TResponse, result)
     finally:
         interrupt.set()

hatchet_sdk/worker/worker.py

@@ -19,13 +19,20 @@
 
 from hatchet_sdk import Context
 from hatchet_sdk.client import Client, new_client_raw
+from hatchet_sdk.clients.dispatcher.action_listener import Action
 from hatchet_sdk.contracts.workflows_pb2 import CreateWorkflowVersionOpts
 from hatchet_sdk.loader import ClientConfig
 from hatchet_sdk.logger import logger
 from hatchet_sdk.utils.types import WorkflowValidator
 from hatchet_sdk.utils.typing import is_basemodel_subclass
-from hatchet_sdk.worker.action_listener_process import worker_action_listener_process
-from hatchet_sdk.worker.runner.run_loop_manager import WorkerActionRunLoopManager
+from hatchet_sdk.worker.action_listener_process import (
+    ActionEvent,
+    worker_action_listener_process,
+)
+from hatchet_sdk.worker.runner.run_loop_manager import (
+    STOP_LOOP_TYPE,
+    WorkerActionRunLoopManager,
+)
 from hatchet_sdk.workflow import WorkflowInterface
 
 T = TypeVar("T")
@@ -74,13 +81,13 @@ def __init__(
         self._status: WorkerStatus
 
         self.action_listener_process: BaseProcess
-        self.action_listener_health_check: asyncio.Task[Any]
+        self.action_listener_health_check: asyncio.Task[None]
         self.action_runner: WorkerActionRunLoopManager
 
         self.ctx = multiprocessing.get_context("spawn")
 
-        self.action_queue: "Queue[Any]" = self.ctx.Queue()
-        self.event_queue: "Queue[Any]" = self.ctx.Queue()
+        self.action_queue: "Queue[Action | STOP_LOOP_TYPE]" = self.ctx.Queue()
+        self.event_queue: "Queue[ActionEvent]" = self.ctx.Queue()
 
         self.loop: asyncio.AbstractEventLoop
 
@@ -193,12 +200,10 @@ async def start_health_server(self) -> None:
 
         logger.info(f"healthcheck server running on port {port}")
 
-    def start(
-        self, options: WorkerStartOptions = WorkerStartOptions()
-    ) -> Future[asyncio.Task[Any] | None]:
+    def start(self, options: WorkerStartOptions = WorkerStartOptions()) -> None:
         self.owned_loop = self.setup_loop(options.loop)
 
-        f = asyncio.run_coroutine_threadsafe(
+        asyncio.run_coroutine_threadsafe(
             self.async_start(options, _from_start=True), self.loop
         )
 
@@ -209,14 +214,12 @@ def start(
             if self.handle_kill:
                 sys.exit(0)
 
-        return f
-
     ## Start methods
     async def async_start(
         self,
         options: WorkerStartOptions = WorkerStartOptions(),
         _from_start: bool = False,
-    ) -> Any | None:
+    ) -> None:
         main_pid = os.getpid()
         logger.info("------------------------------------------")
         logger.info("STARTING HATCHET...")
@@ -245,7 +248,7 @@ async def async_start(
             self._check_listener_health()
         )
 
-        return await self.action_listener_health_check
+        await self.action_listener_health_check
 
     def _run_action_runner(self) -> WorkerActionRunLoopManager:
         # Retrieve the shared queue