Improve types through typing all tests

pyproject.toml

@@ -34,7 +34,7 @@ toml = "^0.10.2"
 mypy = "^1.3.0"
 
 [tool.poetry.group.pyright.dependencies]
-pyright = "^1.1.311"
+pyright = "^1.1.313"
 
 [tool.poetry.group.ruff.dependencies]
 ruff = "^0.0.270"
@@ -48,18 +48,14 @@ pytest-mock = "^3.10.0"
 target-version = ["py38", "py39", "py310", "py311"]
 
 [tool.mypy]
+disallow_untyped_defs = true
 no_implicit_optional = true
 strict_equality = true
 warn_return_any = true
 warn_redundant_casts = true
 warn_unused_ignores = true
 warn_unused_configs = true
 
-[[tool.mypy.overrides]]
-# Options that only applies to src, and not tests
-module = "pykka.*"
-disallow_untyped_defs = true
-
 [tool.ruff]
 select = [
    "A", # flake8-builtins

src/pykka/_actor.py

@@ -4,7 +4,7 @@
 import sys
 import threading
 import uuid
-from typing import TYPE_CHECKING, Any, Optional, Protocol, Sequence
+from typing import TYPE_CHECKING, Any, Optional, Protocol, Sequence, TypeVar
 
 from pykka import ActorDeadError, ActorRef, ActorRegistry, messages
 
@@ -19,6 +19,9 @@
 logger = logging.getLogger("pykka")
 
 
+A = TypeVar("A", bound="Actor")
+
+
 class ActorInbox(Protocol):
     def put(self, envelope: Envelope[Any], /) -> None:
         ...
@@ -75,10 +78,10 @@ def a_method(self, ...):
 
     @classmethod
     def start(
-        cls,
+        cls: type[A],
         *args: Any,
         **kwargs: Any,
-    ) -> ActorRef:
+    ) -> ActorRef[A]:
         """Start an actor.
 
         Starting an actor also registers it in the
@@ -151,8 +154,13 @@ def _start_actor_loop(self) -> None:
     #: friends to put messages in the inbox.
     actor_inbox: ActorInbox
 
-    #: The actor's :class:`ActorRef` instance.
-    actor_ref: ActorRef
+    _actor_ref: ActorRef[Any]
+
+    @property
+    def actor_ref(self: A) -> ActorRef[A]:
+        """The actor's :class:`ActorRef` instance."""
+        # This property only exists to improve the typing of the ActorRef.
+        return self._actor_ref
 
     #: A :class:`threading.Event` representing whether or not the actor should
     #: continue processing messages. Use :meth:`stop` to change it.
@@ -184,7 +192,7 @@ def __init__(
         self.actor_inbox = self._create_actor_inbox()
         self.actor_stopped = threading.Event()
 
-        self.actor_ref = ActorRef(self)
+        self._actor_ref = ActorRef(self)
 
     def __str__(self) -> str:
         return f"{self.__class__.__name__} ({self.actor_urn})"

src/pykka/_future.py

@@ -23,9 +23,8 @@
 
 T = TypeVar("T")
 J = TypeVar("J")  # For when T is Iterable[J]
-
-M = TypeVar("M")  # For Future.map()
-R = TypeVar("R")  # For Future.reduce()
+M = TypeVar("M")  # Result of Future.map()
+R = TypeVar("R")  # Result of Future.reduce()
 
 GetHookFunc: TypeAlias = Callable[[Optional[float]], T]
 
@@ -47,6 +46,9 @@ def __init__(self) -> None:
         self._get_hook = None
         self._get_hook_result = None
 
+    def __repr__(self) -> str:
+        return "<pykka.Future>"
+
     def get(
         self,
         timeout: Optional[float] = None,

src/pykka/_proxy.py

@@ -4,11 +4,11 @@
 from typing import (
     TYPE_CHECKING,
     Any,
+    Generic,
     NamedTuple,
     Optional,
     Sequence,
     TypeVar,
-    Union,
 )
 
 from pykka import ActorDeadError, messages
@@ -23,7 +23,8 @@
 logger = logging.getLogger("pykka")
 
 
-_T = TypeVar("_T")
+T = TypeVar("T")
+A = TypeVar("A", bound="Actor")
 
 AttrPath: TypeAlias = Sequence[str]
 
@@ -33,7 +34,7 @@ class AttrInfo(NamedTuple):
     traversable: bool
 
 
-class ActorProxy:
+class ActorProxy(Generic[A]):
     """An :class:`ActorProxy` wraps an :class:`ActorRef <pykka.ActorRef>` instance.
 
     The proxy allows the referenced actor to be used through regular
@@ -132,17 +133,17 @@ def do_more_work(self):
     """
 
     #: The actor's :class:`pykka.ActorRef` instance.
-    actor_ref: ActorRef
+    actor_ref: ActorRef[A]
 
-    _actor: Actor
+    _actor: A
     _attr_path: AttrPath
     _known_attrs: dict[AttrPath, AttrInfo]
-    _actor_proxies: dict[AttrPath, ActorProxy]
-    _callable_proxies: dict[AttrPath, CallableProxy]
+    _actor_proxies: dict[AttrPath, ActorProxy[A]]
+    _callable_proxies: dict[AttrPath, CallableProxy[A]]
 
     def __init__(
         self,
-        actor_ref: ActorRef,
+        actor_ref: ActorRef[A],
         attr_path: Optional[AttrPath] = None,
     ) -> None:
         if not actor_ref.is_alive():
@@ -229,7 +230,10 @@ def __eq__(
     ) -> bool:
         if not isinstance(other, ActorProxy):
             return False
-        if self._actor != other._actor:  # noqa: SLF001
+        if (
+            self._actor
+            != other._actor  # noqa: SLF001  # pyright: ignore[reportUnknownMemberType]
+        ):
             return False
         if self._attr_path != other._attr_path:  # noqa: SLF001
             return False
@@ -248,10 +252,7 @@ def __dir__(self) -> list[str]:
         result += [attr_path[0] for attr_path in list(self._known_attrs.keys())]
         return sorted(result)
 
-    def __getattr__(
-        self,
-        name: str,
-    ) -> Union[CallableProxy, ActorProxy, Future[Any]]:
+    def __getattr__(self, name: str) -> Any:
         """Get a field or callable from the actor."""
         attr_path: AttrPath = (*self._attr_path, name)
 
@@ -294,7 +295,7 @@ def __setattr__(
         return None
 
 
-class CallableProxy:
+class CallableProxy(Generic[A]):
     """Proxy to a single method.
 
     :class:`CallableProxy` instances are returned when accessing methods on a
@@ -311,12 +312,12 @@ class CallableProxy:
         proxy.do_work.defer()
     """
 
-    actor_ref: ActorRef
+    actor_ref: ActorRef[A]
     _attr_path: AttrPath
 
     def __init__(
         self,
-        actor_ref: ActorRef,
+        actor_ref: ActorRef[A],
         attr_path: AttrPath,
     ) -> None:
         self.actor_ref = actor_ref
@@ -362,7 +363,7 @@ def defer(
         self.actor_ref.tell(message)
 
 
-def traversable(obj: _T) -> _T:
+def traversable(obj: T) -> T:
     """Mark an actor attribute as traversable.
 
     The traversable marker makes the actor attribute's own methods and

src/pykka/_ref.py

@@ -1,6 +1,15 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Literal, Optional, Union, overload
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Generic,
+    Literal,
+    Optional,
+    TypeVar,
+    Union,
+    overload,
+)
 
 from pykka import ActorDeadError, ActorProxy
 from pykka._envelope import Envelope
@@ -9,13 +18,16 @@
 if TYPE_CHECKING:
     from threading import Event
 
-    from pykka import Future
-    from pykka._actor import Actor, ActorInbox
+    from pykka import Actor, Future
+    from pykka._actor import ActorInbox
 
 __all__ = ["ActorRef"]
 
 
-class ActorRef:
+A = TypeVar("A", bound="Actor")
+
+
+class ActorRef(Generic[A]):
     """Reference to a running actor which may safely be passed around.
 
     :class:`ActorRef` instances are returned by :meth:`Actor.start` and the
@@ -27,7 +39,7 @@
     """
 
     #: The class of the referenced actor.
-    actor_class: type[Actor]
+    actor_class: type[A]
 
     #: See :attr:`Actor.actor_urn`.
     actor_urn: str
@@ -39,8 +51,8 @@
     actor_stopped: Event
 
     def __init__(
-        self,
-        actor: Actor,
+        self: ActorRef[A],
+        actor: A,
     ) -> None:
         self._actor = actor
         self.actor_class = actor.__class__
@@ -223,7 +235,7 @@
 
         return converted_future
 
-    def proxy(self) -> ActorProxy:
+    def proxy(self: ActorRef[A]) -> ActorProxy[A]:
         """Wrap the :class:`ActorRef` in an :class:`ActorProxy <pykka.ActorProxy>`.
 
         Using this method like this::

src/pykka/_registry.py

@@ -2,15 +2,25 @@
 
 import logging
 import threading
-from typing import TYPE_CHECKING, Any, ClassVar, Literal, Optional, Union, overload
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    ClassVar,
+    Literal,
+    Optional,
+    TypeVar,
+    Union,
+    overload,
+)
 
 if TYPE_CHECKING:
     from pykka import Actor, ActorRef, Future
 
-logger = logging.getLogger("pykka")
+__all__ = ["ActorRegistry"]
 
+logger = logging.getLogger("pykka")
 
-__all__ = ["ActorRegistry"]
+A = TypeVar("A", bound="Actor")
 
 
 class ActorRegistry:
@@ -19,7 +29,7 @@ class ActorRegistry:
     Contains global state, but should be thread-safe.
     """
 
-    _actor_refs: ClassVar[list[ActorRef]] = []
+    _actor_refs: ClassVar[list[ActorRef[Any]]] = []
     _actor_refs_lock: ClassVar[threading.RLock] = threading.RLock()
 
     @classmethod
@@ -49,7 +59,7 @@ def broadcast(
             ref.tell(message)
 
     @classmethod
-    def get_all(cls) -> list[ActorRef]:
+    def get_all(cls) -> list[ActorRef[Any]]:
         """Get all running actors.
 
         :returns: list of :class:`pykka.ActorRef`
@@ -60,8 +70,8 @@ def get_all(cls) -> list[ActorRef]:
     @classmethod
     def get_by_class(
         cls,
-        actor_class: type[Actor],
-    ) -> list[ActorRef]:
+        actor_class: type[A],
+    ) -> list[ActorRef[A]]:
         """Get all running actors of the given class or a subclass.
 
         :param actor_class: actor class, or any superclass of the actor
@@ -80,7 +90,7 @@ def get_by_class(
     def get_by_class_name(
         cls,
         actor_class_name: str,
-    ) -> list[ActorRef]:
+    ) -> list[ActorRef[Any]]:
         """Get all running actors of the given class name.
 
         :param actor_class_name: actor class name
@@ -99,7 +109,7 @@ def get_by_class_name(
     def get_by_urn(
         cls,
         actor_urn: str,
-    ) -> Optional[ActorRef]:
+    ) -> Optional[ActorRef[Any]]:
         """Get an actor by its universally unique URN.
 
         :param actor_urn: actor URN
@@ -116,7 +126,7 @@ def get_by_urn(
     @classmethod
     def register(
         cls,
-        actor_ref: ActorRef,
+        actor_ref: ActorRef[Any],
     ) -> None:
         """Register an :class:`ActorRef` in the registry.
 
@@ -188,7 +198,7 @@ def stop_all(
     @classmethod
     def unregister(
         cls,
-        actor_ref: ActorRef,
+        actor_ref: ActorRef[A],
     ) -> None:
         """Remove an :class:`ActorRef <pykka.ActorRef>` from the registry.
 

src/pykka/_threading.py

@@ -20,10 +20,10 @@ class ThreadingFutureResult(NamedTuple):
     exc_info: Optional[OptExcInfo] = None
 
 
-F = TypeVar("F")
+T = TypeVar("T")
 
 
-class ThreadingFuture(Future[F]):
+class ThreadingFuture(Future[T]):
     """Implementation of :class:`Future` for use with regular Python threads`.
 
     The future is implemented using a :class:`queue.Queue`.