typing: propagate result types through async APIs

kazoo/client.py

@@ -59,6 +59,7 @@
     SetData,
     Sync,
     Transaction,
+    Transaction_Response,
 )
 from kazoo.protocol.states import (
     Callback,
@@ -1038,9 +1039,11 @@ def add_auth(self, scheme: str, credential: str) -> bool:
             the session state will be set to AUTH_FAILED as well.
 
         """
-        return cast("bool", self.add_auth_async(scheme, credential).get())
+        return self.add_auth_async(scheme, credential).get()
 
-    def add_auth_async(self, scheme: str, credential: str) -> IAsyncResult:
+    def add_auth_async(
+        self, scheme: str, credential: str
+    ) -> IAsyncResult[bool]:
         """Asynchronously send credentials to server. Takes the same
         arguments as :meth:`add_auth`.
 
@@ -1070,7 +1073,7 @@ def unchroot(self, path: str) -> str:
         else:
             return path
 
-    def sync_async(self, path: str) -> IAsyncResult:
+    def sync_async(self, path: str) -> IAsyncResult[str]:
         """Asynchronous sync.
 
         :rtype: :class:`~kazoo.interfaces.IAsyncResult`
@@ -1104,7 +1107,7 @@ def sync(self, path: str) -> str:
         .. versionadded:: 0.5
 
         """
-        return cast("str", self.sync_async(path).get())
+        return self.sync_async(path).get()
 
     @overload
     def create(
@@ -1220,18 +1223,15 @@ def create(
             The `include_data` option.
         """
         acl = acl or self.default_acl
-        return cast(
-            "str | tuple[str, ZnodeStat]",
-            self.create_async(
-                path,
-                value,
-                acl=acl,
-                ephemeral=ephemeral,
-                sequence=sequence,
-                makepath=makepath,
-                include_data=include_data,
-            ).get(),
-        )
+        return self.create_async(
+            path,
+            value,
+            acl=acl,
+            ephemeral=ephemeral,
+            sequence=sequence,
+            makepath=makepath,
+            include_data=include_data,
+        ).get()
 
     def create_async(
         self,
@@ -1242,7 +1242,7 @@ def create_async(
         sequence: bool = False,
         makepath: bool = False,
         include_data: bool = False,
-    ) -> IAsyncResult:
+    ) -> IAsyncResult[str | tuple[str, ZnodeStat]]:
         """Asynchronously create a ZNode. Takes the same arguments as
         :meth:`create`.
 
@@ -1338,7 +1338,7 @@ def _create_async_inner(
         flags: int,
         trailing: bool = False,
         include_data: bool = False,
-    ) -> IAsyncResult:
+    ) -> IAsyncResult[str | tuple[str, ZnodeStat]]:
         async_result = self.handler.async_result()
         opcode = Create2 if include_data else Create
 
@@ -1370,11 +1370,11 @@ def ensure_path(self, path: str, acl: Sequence[ACL] | None = None) -> bool:
         :param acl: Permissions for node.
 
         """
-        return cast("bool", self.ensure_path_async(path, acl).get())
+        return self.ensure_path_async(path, acl).get()
 
     def ensure_path_async(
         self, path: str, acl: Sequence[ACL] | None = None
-    ) -> IAsyncResult:
+    ) -> IAsyncResult[bool]:
         """Recursively create a path asynchronously if it doesn't
         exist. Takes the same arguments as :meth:`ensure_path`.
 
@@ -1439,13 +1439,11 @@ def exists(
             returns a non-zero error code.
 
         """
-        return cast(
-            "ZnodeStat | None", self.exists_async(path, watch=watch).get()
-        )
+        return self.exists_async(path, watch=watch).get()
 
     def exists_async(
         self, path: str, watch: WatchFunc | None = None
-    ) -> IAsyncResult:
+    ) -> IAsyncResult[ZnodeStat | None]:
         """Asynchronously check if a node exists. Takes the same
         arguments as :meth:`exists`.
 
@@ -1488,13 +1486,11 @@ def get(
             returns a non-zero error code
 
         """
-        return cast(
-            "tuple[bytes, ZnodeStat]", self.get_async(path, watch=watch).get()
-        )
+        return self.get_async(path, watch=watch).get()
 
     def get_async(
         self, path: str, watch: WatchFunc | None = None
-    ) -> IAsyncResult:
+    ) -> IAsyncResult[tuple[bytes, ZnodeStat]]:
         """Asynchronously get the value of a node. Takes the same
         arguments as :meth:`get`.
 
@@ -1569,19 +1565,16 @@ def get_children(
             The `include_data` option.
 
         """
-        return cast(
-            "list[str] | tuple[list[str], ZnodeStat]",
-            self.get_children_async(
-                path, watch=watch, include_data=include_data
-            ).get(),
-        )
+        return self.get_children_async(
+            path, watch=watch, include_data=include_data
+        ).get()
 
     def get_children_async(
         self,
         path: str,
         watch: WatchFunc | None = None,
         include_data: bool = False,
-    ) -> IAsyncResult:
+    ) -> IAsyncResult[list[str] | tuple[list[str], ZnodeStat]]:
         """Asynchronously get a list of child nodes of a path. Takes
         the same arguments as :meth:`get_children`.
 
@@ -1623,11 +1616,11 @@ def get_acls(self, path: str) -> tuple[list[ACL], ZnodeStat]:
         .. versionadded:: 0.5
 
         """
-        return cast(
-            "tuple[list[ACL], ZnodeStat]", self.get_acls_async(path).get()
-        )
+        return self.get_acls_async(path).get()
 
-    def get_acls_async(self, path: str) -> IAsyncResult:
+    def get_acls_async(
+        self, path: str
+    ) -> IAsyncResult[tuple[list[ACL], ZnodeStat]]:
         """Return the ACL and stat of the node of the given path. Takes
         the same arguments as :meth:`get_acls`.
 
@@ -1670,13 +1663,11 @@ def set_acls(
         .. versionadded:: 0.5
 
         """
-        return cast(
-            "ZnodeStat", self.set_acls_async(path, acls, version).get()
-        )
+        return self.set_acls_async(path, acls, version).get()
 
     def set_acls_async(
         self, path: str, acls: Sequence[ACL], version: int = -1
-    ) -> IAsyncResult:
+    ) -> IAsyncResult[ZnodeStat]:
         """Set the ACL for the node of the given path. Takes the same
         arguments as :meth:`set_acls`.
 
@@ -1734,11 +1725,11 @@ def set(
             returns a non-zero error code.
 
         """
-        return cast("ZnodeStat", self.set_async(path, value, version).get())
+        return self.set_async(path, value, version).get()
 
     def set_async(
         self, path: str, value: bytes | None, version: int = -1
-    ) -> IAsyncResult:
+    ) -> IAsyncResult[ZnodeStat]:
         """Set the value of a node. Takes the same arguments as
         :meth:`set`.
 
@@ -1819,7 +1810,7 @@ def delete(
         else:
             return self.delete_async(path, version).get()
 
-    def delete_async(self, path: str, version: int = -1) -> IAsyncResult:
+    def delete_async(self, path: str, version: int = -1) -> IAsyncResult[bool]:
         """Asynchronously delete a node. Takes the same arguments as
         :meth:`delete`, with the exception of `recursive`.
 
@@ -1933,15 +1924,15 @@ def reconfig(
         result = self.reconfig_async(
             joining, leaving, new_members, from_config
         )
-        return cast("tuple[bytes, ZnodeStat]", result.get())
+        return result.get()
 
     def reconfig_async(
         self,
         joining: str | None,
         leaving: str | None,
         new_members: str | None,
         from_config: int,
-    ) -> IAsyncResult:
+    ) -> IAsyncResult[tuple[bytes, ZnodeStat]]:
         """Asynchronously reconfig a cluster. Takes the same arguments as
         :meth:`reconfig`.
 
@@ -2079,7 +2070,7 @@ def check(self, path: str, version: int) -> None:
             CheckVersion(_prefix_root(self.client.chroot, path), version)
         )
 
-    def commit_async(self) -> IAsyncResult:
+    def commit_async(self) -> IAsyncResult[list[Transaction_Response]]:
         """Commit the transaction asynchronously.
 
         :rtype: :class:`~kazoo.interfaces.IAsyncResult`
@@ -2091,14 +2082,14 @@ def commit_async(self) -> IAsyncResult:
         self.client._call(Transaction(self.operations), async_object)
         return async_object
 
-    def commit(self) -> list[Any]:
+    def commit(self) -> list[Transaction_Response]:
         """Commit the transaction.
 
         :returns: A list of the results for each operation in the
                   transaction.
 
         """
-        return cast("list[Any]", self.commit_async().get())
+        return self.commit_async().get()
 
     def __enter__(self) -> TransactionRequest:
         return self

kazoo/handlers/utils.py

@@ -33,7 +33,7 @@
 # want to change the code too much.
 _NONE = object()
 
-CallbackFunc = Callable[..., None]
+CallbackFunc = Callable[..., object]
 
 
 class AsyncResult(IAsyncResult):

kazoo/interfaces.py

@@ -17,8 +17,10 @@
     Iterable,
     Protocol,
     Union,
+    overload,
     TYPE_CHECKING,
 )
+from typing_extensions import TypeVar
 
 if TYPE_CHECKING:
     from types import TracebackType
@@ -158,6 +160,7 @@ def join(self, timeout: float | None = None) -> None:
 
 
 SpawnedFunc = Callable[..., None]
+_ResultT = TypeVar("_ResultT", default=Any)
 
 
 class IHandler(Protocol):
@@ -239,10 +242,8 @@ def rlock_object(self) -> ReentrantLock:
         """Return an appropriate object that implements Python's
         threading.RLock API"""
 
-    def async_result(self) -> IAsyncResult:
-        """Return an instance that conforms to the
-        :class:`~IAsyncResult` interface appropriate for this
-        handler"""
+    def async_result(self) -> IAsyncResult[_ResultT]:
+        """Return a typed asynchronous result for this handler."""
 
     def spawn(
         self, func: SpawnedFunc, *args: Any, **kwargs: Any
@@ -266,7 +267,7 @@ def dispatch_callback(self, callback: Callback) -> None:
         """
 
 
-class IAsyncResult(Protocol):
+class IAsyncResult(Protocol[_ResultT]):
     """An Async Result object that can be queried for a value that has
     been set asynchronously.
 
@@ -297,7 +298,12 @@ def successful(self) -> bool:
         """Return `True` if and only if it is ready and holds a
         value"""
 
-    def set(self, value: Any = None) -> None:
+    @overload
+    def set(self: IAsyncResult[None]) -> None:
+        ...
+
+    @overload
+    def set(self, value: _ResultT) -> None:
         """Store the value. Wake up the waiters.
 
         :param value: Value to store as the result.
@@ -315,7 +321,9 @@ def set_exception(self, exception: Exception) -> None:
         up. Sequential calls to :meth:`wait` and :meth:`get` will not
         block at all."""
 
-    def get(self, block: bool = True, timeout: float | None = None) -> Any:
+    def get(
+        self, block: bool = True, timeout: float | None = None
+    ) -> _ResultT:
         """Return the stored value or raise the exception
 
         :param block: Whether this method should block or return
@@ -330,15 +338,18 @@ def get(self, block: bool = True, timeout: float | None = None) -> Any:
         :meth:`set_exception` has been called or until the optional
         timeout occurs."""
 
-    def get_nowait(self) -> Any:
+    def get_nowait(self) -> _ResultT:
         """Return the value or raise the exception without blocking.
 
         If nothing is available, raise the Timeout exception class on
         the associated :class:`IHandler` interface."""
 
-    def wait(self, timeout: float | None = None) -> Any:
+    def wait(self, timeout: float | None = None) -> object:
         """Block until the instance is ready.
 
+        Handler implementations differ in what ``wait`` returns; use
+        :meth:`get` when the typed result value is needed.
+
         :param timeout: How long to wait for a value when `block` is
                         `True`.
         :type timeout: float
@@ -348,7 +359,10 @@ def wait(self, timeout: float | None = None) -> Any:
         :meth:`set_exception` has been called or until the optional
         timeout occurs."""
 
-    def rawlink(self, callback: Callable[[IAsyncResult], Any]) -> None:
+    def rawlink(
+        self,
+        callback: Callable[[IAsyncResult[_ResultT]], object],
+    ) -> None:
         """Register a callback to call when a value or an exception is
         set
 
@@ -360,7 +374,10 @@ def rawlink(self, callback: Callable[[IAsyncResult], Any]) -> None:
 
         """
 
-    def unlink(self, callback: Callable[[IAsyncResult], None]) -> None:
+    def unlink(
+        self,
+        callback: Callable[[IAsyncResult[_ResultT]], object],
+    ) -> None:
         """Remove the callback set by :meth:`rawlink`
 
         :param callback: A callback function to remove.

kazoo/tests/test_client.py

@@ -1202,6 +1202,7 @@ def test_basic_create(self) -> None:
         results = t.commit()
         assert len(results) == 3
         assert results[0] == "/freddy"
+        assert isinstance(results[2], str)
         assert results[2].startswith("/smith0") is True
 
     def test_bad_creates(self) -> None: