Merge pull request #804 from LincolnPuzey/more_actor_type_fixes

Fixes to type hints for `@actor` decorator and `Actor` class

Author: LincolnPuzey

dramatiq/actor.py

@@ -63,7 +63,7 @@ class Actor(Generic[P, R]):
 
     def __init__(
         self,
-        fn: Callable[P, Union[R, Awaitable[R]]],
+        fn: Union[Callable[P, Awaitable[R]], Callable[P, R]],
         *,
         broker: Broker,
         actor_name: str,
@@ -75,7 +75,7 @@ def __init__(
             raise ValueError(f"An actor named {actor_name!r} is already registered.")
 
         self.logger = get_logger(fn.__module__ or "_", actor_name)
-        self.fn = async_to_sync(fn) if iscoroutinefunction(fn) else fn
+        self.fn: Callable[P, R] = async_to_sync(fn) if iscoroutinefunction(fn) else fn  # type: ignore[assignment]
         self.broker = broker
         self.actor_name = actor_name
         self.queue_name = queue_name
@@ -178,7 +178,7 @@ def send_with_options(
         message = self.message_with_options(args=args, kwargs=kwargs, **options)
         return self.broker.enqueue(message, delay=delay)
 
-    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> Any | R | Awaitable[R]:
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> R:
         """Synchronously call this actor.
 
         Parameters:
@@ -210,21 +210,29 @@ def __call__(self, fn: Callable[P, Awaitable[R]]) -> Actor[P, R]: ...
     @overload
     def __call__(self, fn: Callable[P, R]) -> Actor[P, R]: ...
 
-    def __call__(self, fn: Callable[P, Union[Awaitable[R], R]]) -> Actor[P, R]: ...
+    def __call__(self, fn: Union[Callable[P, Awaitable[R]], Callable[P, R]]) -> Actor[P, R]: ...
 
 
 @overload
-def actor(fn: Callable[P, Awaitable[R]], **kwargs) -> Actor[P, R]:
+def actor(fn: Callable[P, Awaitable[R]]) -> Actor[P, R]:
     pass
 
 
 @overload
-def actor(fn: Callable[P, R], **kwargs) -> Actor[P, R]:
+def actor(fn: Callable[P, R]) -> Actor[P, R]:
     pass
 
 
 @overload
-def actor(fn: None = None, **kwargs) -> ActorDecorator:
+def actor(
+    *,
+    queue_name: str = "default",
+    priority: int = 0,
+    actor_name: Optional[str] = None,
+    broker: Optional[Broker] = None,
+    actor_class: Callable[..., Actor[Any, Any]] = Actor,
+    **options: Any,
+) -> ActorDecorator:
     pass
 
 
@@ -236,7 +244,7 @@ def actor(
     queue_name: str = "default",
     priority: int = 0,
     broker: Optional[Broker] = None,
-    **options,
+    **options: Any,
 ) -> Union[Actor[P, R], Callable]:
     """Declare an actor.
 

dramatiq/asyncio.py

@@ -22,7 +22,7 @@
 import functools
 import logging
 import threading
-from typing import Awaitable, Callable, Optional, TypeVar
+from typing import Awaitable, Callable, Optional, ParamSpec, TypeVar
 
 from .threading import Interrupt
 
@@ -34,6 +34,7 @@
 ]
 
 R = TypeVar("R")
+P = ParamSpec("P")
 
 _event_loop_thread = None
 
@@ -53,13 +54,13 @@ def set_event_loop_thread(thread: Optional[EventLoopThread]) -> None:
     _event_loop_thread = thread
 
 
-def async_to_sync(async_fn: Callable[..., Awaitable[R]]) -> Callable[..., R]:
+def async_to_sync(async_fn: Callable[P, Awaitable[R]]) -> Callable[P, R]:
     """Wrap an async function to run it on the event loop thread and
     synchronously wait for its result on the calling thread.
     """
 
     @functools.wraps(async_fn)
-    def wrapper(*args, **kwargs) -> R:
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
         event_loop_thread = get_event_loop_thread()
         if event_loop_thread is None:
             raise RuntimeError(

tests/test_types.py

@@ -0,0 +1,114 @@
+"""Tests for dramatiq's types.
+
+Unlike other test files which are run with pytest,
+this test file should be type-checked with mypy,
+to test that Dramatiq's types can be "consumed" by user code without type errors.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, ParamSpec, TypeVar
+
+if TYPE_CHECKING:
+    from typing_extensions import assert_type
+
+import dramatiq
+
+P = ParamSpec("P")
+R = TypeVar("R")
+
+broker = dramatiq.get_broker()
+
+
+class ArgType:
+    pass
+
+
+class ReturnType:
+    pass
+
+
+class CustomActor(dramatiq.Actor[P, R]):
+    pass
+
+
+# # # Tests for @actor decorator # # #
+
+
+@dramatiq.actor
+def actor(arg: ArgType) -> ReturnType:
+    return ReturnType()
+
+
+@dramatiq.actor()
+def actor_no_options(arg: ArgType) -> ReturnType:
+    return ReturnType()
+
+
+@dramatiq.actor(
+    actor_name="actor_with_options",
+    queue_name="some_queue",
+    priority=2,
+    broker=broker,
+    max_age=1,
+)
+def actor_with_options(arg: ArgType) -> ReturnType:
+    return ReturnType()
+
+
+@dramatiq.actor(
+    actor_class=CustomActor,
+    actor_name="actor_with_custom_actor_class",
+    queue_name="some_queue",
+    priority=2,
+    broker=broker,
+    max_age=1,
+)
+def actor_with_custom_actor_class(arg: ArgType) -> ReturnType:
+    return ReturnType()
+
+
+@dramatiq.actor
+async def async_actor(arg: ArgType) -> ReturnType:
+    return ReturnType()
+
+
+@dramatiq.actor()
+async def async_actor_no_options(arg: ArgType) -> ReturnType:
+    return ReturnType()
+
+
+@dramatiq.actor(
+    actor_name="async_actor_with_options",
+    queue_name="some_queue",
+    priority=2,
+    broker=broker,
+    max_age=1,
+)
+async def async_actor_with_options(arg: ArgType) -> ReturnType:
+    return ReturnType()
+
+
+@dramatiq.actor(
+    actor_class=CustomActor,
+    actor_name="async_actor_with_custom_actor_class",
+    queue_name="some_queue",
+    priority=2,
+    broker=broker,
+    max_age=1,
+)
+async def async_actor_with_custom_actor_class(arg: ArgType) -> ReturnType:
+    return ReturnType()
+
+
+# # # Test that calling actors has correct arg/return type
+def _calling_actors_type_check() -> None:
+    assert_type(actor(ArgType()), ReturnType)
+    assert_type(actor_no_options(ArgType()), ReturnType)
+    assert_type(actor_with_options(ArgType()), ReturnType)
+    assert_type(actor_with_custom_actor_class(ArgType()), ReturnType)
+
+    assert_type(async_actor(ArgType()), ReturnType)
+    assert_type(async_actor_no_options(ArgType()), ReturnType)
+    assert_type(async_actor_with_options(ArgType()), ReturnType)
+    assert_type(async_actor_with_custom_actor_class(ArgType()), ReturnType)