fix: proper typing for call and acall of Step

Author: mrkaye97

examples/simple/worker.py

@@ -1,10 +1,10 @@
 import time
 
+from dotenv import load_dotenv
+
 from hatchet_sdk import Context
 from hatchet_sdk.v2 import Hatchet, Workflow, WorkflowConfig
 
-from dotenv import load_dotenv
-
 load_dotenv()
 
 hatchet = Hatchet(debug=True)

hatchet_sdk/v2/workflows.py

@@ -1,6 +1,18 @@
 import asyncio
 from enum import Enum
-from typing import Any, Callable, Generic, ParamSpec, Type, TypeVar, Union, cast
+from typing import (
+    Any,
+    Awaitable,
+    Callable,
+    Generic,
+    ParamSpec,
+    Type,
+    TypeGuard,
+    TypeVar,
+    Union,
+    cast,
+    overload,
+)
 
 from pydantic import BaseModel, ConfigDict
 
@@ -16,8 +28,6 @@
 )
 from hatchet_sdk.contracts.workflows_pb2 import StickyStrategy as StickyStrategyProto
 from hatchet_sdk.contracts.workflows_pb2 import WorkflowConcurrencyOpts, WorkflowKind
-from hatchet_sdk.labels import DesiredWorkerLabel
-from hatchet_sdk.rate_limit import RateLimit
 
 from ..logger import logger
 
@@ -57,6 +67,7 @@ class StickyStrategy(str, Enum):
 
 class WorkflowConfig(BaseModel):
     model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)
+
     name: str = ""
     on_events: list[str] = []
     on_crons: list[str] = []
@@ -75,10 +86,23 @@ class StepType(str, Enum):
     ON_FAILURE = "on_failure"
 
 
+AsyncFunc = Callable[[Any, Context], Awaitable[R]]
+SyncFunc = Callable[[Any, Context], R]
+StepFunc = Union[AsyncFunc[R], SyncFunc[R]]
+
+
+def is_async_fn(fn: StepFunc[R]) -> TypeGuard[AsyncFunc[R]]:
+    return asyncio.iscoroutinefunction(fn)
+
+
+def is_sync_fn(fn: StepFunc[R]) -> TypeGuard[SyncFunc[R]]:
+    return not asyncio.iscoroutinefunction(fn)
+
+
 class Step(Generic[R]):
     def __init__(
         self,
-        fn: Callable[[Any, Context], R],
+        fn: Callable[[Any, Context], R] | Callable[[Any, Context], Awaitable[R]],
         type: StepType,
         name: str = "",
         timeout: str = "60m",
@@ -90,7 +114,7 @@ def __init__(
         backoff_max_seconds: int | None = None,
     ) -> None:
         self.fn = fn
-        self.is_async_function = bool(asyncio.iscoroutinefunction(fn))
+        self.is_async_function = is_async_fn(fn)
 
         self.type = type
         self.timeout = timeout
@@ -104,8 +128,28 @@ def __init__(
         self.concurrency__max_runs = 1
         self.concurrency__limit_strategy = ConcurrencyLimitStrategy.CANCEL_IN_PROGRESS
 
-    def __call__(self, ctx: Context) -> R:
-        return self.fn(None, ctx)
+    def call(self, ctx: Context) -> R:
+        if self.is_async_function:
+            raise TypeError(f"{self.name} is not a sync function. Use `acall` instead.")
+
+        sync_fn = self.fn
+        if is_sync_fn(sync_fn):
+            return sync_fn(None, ctx)
+
+        raise TypeError(f"{self.name} is not a sync function. Use `acall` instead.")
+
+    async def acall(self, ctx: Context) -> R:
+        if not self.is_async_function:
+            raise TypeError(
+                f"{self.name} is not an async function. Use `call` instead."
+            )
+
+        async_fn = self.fn
+
+        if is_async_fn(async_fn):
+            return await async_fn(None, ctx)
+
+        raise TypeError(f"{self.name} is not an async function. Use `call` instead.")
 
 
 class Workflow:

hatchet_sdk/worker/runner/runner.py

@@ -47,6 +47,8 @@
 from hatchet_sdk.worker.action_listener_process import ActionEvent
 from hatchet_sdk.worker.runner.utils.capture_logs import copy_context_vars, sr, wr
 
+T = TypeVar("T")
+
 if TYPE_CHECKING:
     from hatchet_sdk.v2.workflows import Step
 
@@ -65,7 +67,7 @@ def __init__(
         event_queue: "Queue[Any]",
         max_runs: int | None = None,
         handle_kill: bool = True,
-        action_registry: dict[str, "Step[Any]"] = {},
+        action_registry: dict[str, "Step[T]"] = {},
         validator_registry: dict[str, WorkflowValidator] = {},
         config: ClientConfig = ClientConfig(),
         labels: dict[str, str | int] = {},
@@ -77,7 +79,7 @@ def __init__(
         self.max_runs = max_runs
         self.tasks: dict[str, asyncio.Task[Any]] = {}  # Store run ids and futures
         self.contexts: dict[str, Context] = {}  # Store run ids and contexts
-        self.action_registry: dict[str, "Step[Any]"] = action_registry
+        self.action_registry: dict[str, "Step[T]"] = action_registry
         self.validator_registry = validator_registry
 
         self.event_queue = event_queue
@@ -223,9 +225,7 @@ def inner_callback(task: asyncio.Task[Any]) -> None:
         return inner_callback
 
     ## TODO: Stricter type hinting here
-    def thread_action_func(
-        self, context: Context, action_func: Callable[..., Any], action: Action
-    ) -> Any:
+    def thread_action_func(self, context: Context, step: Step[T], action: Action) -> T:
         if action.step_run_id is not None and action.step_run_id != "":
             self.threads[action.step_run_id] = current_thread()
         elif (
@@ -234,25 +234,23 @@ def thread_action_func(
         ):
             self.threads[action.get_group_key_run_id] = current_thread()
 
-        return action_func(context)
+        return step.call(context)
 
     ## TODO: Stricter type hinting here
     # We wrap all actions in an async func
     async def async_wrapped_action_func(
         self,
         context: Context,
-        action_func: Callable[..., Any],
+        step: Step[T],
         action: Action,
         run_id: str,
-    ) -> Any:
+    ) -> T:
         wr.set(context.workflow_run_id())
         sr.set(context.step_run_id)
 
         try:
-            if (
-                hasattr(action_func, "is_coroutine") and action_func.is_coroutine
-            ) or asyncio.iscoroutinefunction(action_func):
-                return await action_func(context)
+            if step.is_async_function:
+                return await step.acall(context)
             else:
                 pfunc = functools.partial(
                     # we must copy the context vars to the new thread, as only asyncio natively supports
@@ -261,7 +259,7 @@ async def async_wrapped_action_func(
                     contextvars.copy_context().items(),
                     self.thread_action_func,
                     context,
-                    action_func,
+                    step,
                     action,
                 )