From 48738be725ce13609a42fae335ab3fe97f823fbd Mon Sep 17 00:00:00 2001 From: Jon Bringhurst Date: Mon, 31 Aug 2026 18:48:35 -0700 Subject: [PATCH] typing: propagate result types through async APIs ## Why is this needed? Kazoo's async methods return a bare IAsyncResult, so get() becomes Any and synchronous wrappers need casts. This makes the public annotations less useful. This is the first step toward typing the full request path. A follow-up can tie each wire request to its response with a generic _submit helper, then carry that type through the pending queue. ## Proposed Changes - Make IAsyncResult generic in the value it carries. - Add result types to the public async client methods. - Remove casts that are no longer needed. - Type ignored callback and handler-specific wait() returns as object. - Add static checks for the public API. ## Does this PR introduce any breaking change? No. Runtime APIs and behavior stay the same. --- kazoo/client.py | 93 ++++++++++++++++-------------------- kazoo/handlers/utils.py | 2 +- kazoo/interfaces.py | 39 ++++++++++----- kazoo/tests/test_client.py | 1 + kazoo/tests/typing_client.py | 48 +++++++++++++++++++ 5 files changed, 120 insertions(+), 63 deletions(-) create mode 100644 kazoo/tests/typing_client.py diff --git a/kazoo/client.py b/kazoo/client.py index 97bef453..482559cd 100644 --- a/kazoo/client.py +++ b/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,7 +1924,7 @@ 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, @@ -1941,7 +1932,7 @@ def reconfig_async( 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 diff --git a/kazoo/handlers/utils.py b/kazoo/handlers/utils.py index 7ea5d4f8..b1b00e6c 100644 --- a/kazoo/handlers/utils.py +++ b/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): diff --git a/kazoo/interfaces.py b/kazoo/interfaces.py index 466067e3..ec1bc7ae 100644 --- a/kazoo/interfaces.py +++ b/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. diff --git a/kazoo/tests/test_client.py b/kazoo/tests/test_client.py index a031b1ff..39967236 100644 --- a/kazoo/tests/test_client.py +++ b/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: diff --git a/kazoo/tests/typing_client.py b/kazoo/tests/typing_client.py new file mode 100644 index 00000000..9072001e --- /dev/null +++ b/kazoo/tests/typing_client.py @@ -0,0 +1,48 @@ +"""Static checks for the public type annotations.""" + +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from typing_extensions import assert_type + + from kazoo.client import KazooClient + from kazoo.interfaces import IAsyncResult + from kazoo.protocol.serialization import Transaction_Response + from kazoo.protocol.states import ZnodeStat + from kazoo.security import ACL + + client = KazooClient() + + assert_type( + client.add_auth_async("digest", "user:pass"), IAsyncResult[bool] + ) + assert_type(client.sync_async("/node"), IAsyncResult[str]) + assert_type( + client.create_async("/node"), + IAsyncResult[str | tuple[str, ZnodeStat]], + ) + assert_type(client.ensure_path_async("/node"), IAsyncResult[bool]) + assert_type(client.exists_async("/node"), IAsyncResult[ZnodeStat | None]) + assert_type( + client.get_async("/node"), + IAsyncResult[tuple[bytes, ZnodeStat]], + ) + assert_type( + client.get_children_async("/node"), + IAsyncResult[list[str] | tuple[list[str], ZnodeStat]], + ) + assert_type( + client.get_acls_async("/node"), + IAsyncResult[tuple[list[ACL], ZnodeStat]], + ) + assert_type(client.set_acls_async("/node", []), IAsyncResult[ZnodeStat]) + assert_type(client.set_async("/node", b"data"), IAsyncResult[ZnodeStat]) + assert_type(client.delete_async("/node"), IAsyncResult[bool]) + assert_type( + client.reconfig_async(None, None, None, -1), + IAsyncResult[tuple[bytes, ZnodeStat]], + ) + assert_type( + client.transaction().commit_async(), + IAsyncResult[list[Transaction_Response]], + )