docs: add types to recipe.party

Author: a-ungurianu

kazoo/client.py

@@ -68,6 +68,7 @@
 
 if TYPE_CHECKING:
     from typing import (
+        Any,
         List,
         Optional,
         Sequence,
@@ -284,17 +285,17 @@ def __init__(
         self._stopped.set()
         self._writer_stopped.set()
 
-        self.retry = self._conn_retry = None
+        _retry = self._conn_retry = None
 
         if type(connection_retry) is dict:
             self._conn_retry = KazooRetry(**connection_retry)
         elif type(connection_retry) is KazooRetry:
             self._conn_retry = connection_retry
 
         if type(command_retry) is dict:
-            self.retry = KazooRetry(**command_retry)
+            _retry = KazooRetry(**command_retry)
         elif type(command_retry) is KazooRetry:
-            self.retry = command_retry
+            _retry = command_retry
 
         if type(self._conn_retry) is KazooRetry:
             if self.handler.sleep_func != self._conn_retry.sleep_func:
@@ -303,14 +304,14 @@ def __init__(
                     " must use the same sleep func"
                 )
 
-        if type(self.retry) is KazooRetry:
-            if self.handler.sleep_func != self.retry.sleep_func:
+        if type(_retry) is KazooRetry:
+            if self.handler.sleep_func != _retry.sleep_func:
                 raise ConfigurationError(
                     "Command retry handler and event handler "
                     "must use the same sleep func"
                 )
 
-        if self.retry is None or self._conn_retry is None:
+        if _retry is None or self._conn_retry is None:
             old_retry_keys = dict(_RETRY_COMPAT_DEFAULTS)
             for key in old_retry_keys:
                 try:
@@ -326,16 +327,16 @@ def __init__(
                 except KeyError:
                     pass
 
-            retry_keys = {}
+            retry_keys: Any = {}
             for oldname, value in old_retry_keys.items():
                 retry_keys[_RETRY_COMPAT_MAPPING[oldname]] = value
 
             if self._conn_retry is None:
                 self._conn_retry = KazooRetry(
                     sleep_func=self.handler.sleep_func, **retry_keys
                 )
-            if self.retry is None:
-                self.retry = KazooRetry(
+            if _retry is None:
+                _retry = KazooRetry(
                     sleep_func=self.handler.sleep_func, **retry_keys
                 )
 
@@ -380,14 +381,7 @@ def __init__(
             sasl_options=sasl_options,
         )
 
-        # Every retry call should have its own copy of the retry helper
-        # to avoid shared retry counts
-        self._retry = self.retry
-
-        def _retry(*args, **kwargs):
-            return self._retry.copy()(*args, **kwargs)
-
-        self.retry = _retry
+        self._retry = _retry
 
         self.Barrier = partial(Barrier, self)
         self.Counter = partial(Counter, self)
@@ -414,6 +408,12 @@ def _retry(*args, **kwargs):
                 % (kwargs.keys(),)
             )
 
+    @property
+    def retry(self) -> KazooRetry:
+        # Every retry call should have its own copy of the retry helper
+        # to avoid shared retry counts
+        return self._retry.copy()
+
     def _reset(self):
         """Resets a variety of client states for a new connection."""
         self._queue = deque()

kazoo/recipe/party.py

@@ -7,15 +7,26 @@
 used for determining members of a party.
 
 """
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING
 import uuid
 
 from kazoo.exceptions import NodeExistsError, NoNodeError
 
+if TYPE_CHECKING:
+    from typing import Generator, List, Optional
+
+    from kazoo.client import KazooClient
+
 
-class BaseParty(object):
+class BaseParty(ABC):
     """Base implementation of a party."""
 
-    def __init__(self, client, path, identifier=None):
+    def __init__(
+        self, client: KazooClient, path: str, identifier: Optional[str] = None
+    ):
         """
         :param client: A :class:`~kazoo.client.KazooClient` instance.
         :param path: The party path to use.
@@ -29,17 +40,17 @@ def __init__(self, client, path, identifier=None):
         self.ensured_path = False
         self.participating = False
 
-    def _ensure_parent(self):
+    def _ensure_parent(self) -> None:
         if not self.ensured_path:
             # make sure our parent node exists
             self.client.ensure_path(self.path)
             self.ensured_path = True
 
-    def join(self):
+    def join(self) -> None:
         """Join the party"""
         return self.client.retry(self._inner_join)
 
-    def _inner_join(self):
+    def _inner_join(self) -> None:
         self._ensure_parent()
         try:
             self.client.create(self.create_path, self.data, ephemeral=True)
@@ -49,38 +60,49 @@ def _inner_join(self):
             # suspended connection
             self.participating = True
 
-    def leave(self):
+    def leave(self) -> bool:
         """Leave the party"""
         self.participating = False
         return self.client.retry(self._inner_leave)
 
-    def _inner_leave(self):
+    def _inner_leave(self) -> bool:
         try:
             self.client.delete(self.create_path)
         except NoNodeError:
             return False
         return True
 
-    def __len__(self):
+    def __len__(self) -> int:
         """Return a count of participating clients"""
         self._ensure_parent()
         return len(self._get_children())
 
-    def _get_children(self):
+    def _get_children(self) -> List[str]:
         return self.client.retry(self.client.get_children, self.path)
 
+    @property
+    @abstractmethod
+    def create_path(self) -> str:
+        ...
+
 
 class Party(BaseParty):
     """Simple pool of participating processes"""
 
     _NODE_NAME = "__party__"
 
-    def __init__(self, client, path, identifier=None):
+    def __init__(
+        self, client: KazooClient, path: str, identifier: Optional[str] = None
+    ):
         BaseParty.__init__(self, client, path, identifier=identifier)
         self.node = uuid.uuid4().hex + self._NODE_NAME
-        self.create_path = self.path + "/" + self.node
+        self._create_path = self.path + "/" + self.node
+
+    @property
+    def create_path(self) -> str:
+        return self._create_path
 
-    def __iter__(self):
+    def __iter__(self) -> Generator[str, None, None]:
         """Get a list of participating clients' data values"""
         self._ensure_parent()
         children = self._get_children()
@@ -93,7 +115,7 @@ def __iter__(self):
             except NoNodeError:  # pragma: nocover
                 pass
 
-    def _get_children(self):
+    def _get_children(self) -> List[str]:
         children = BaseParty._get_children(self)
         return [c for c in children if self._NODE_NAME in c]
 
@@ -109,12 +131,18 @@ class ShallowParty(BaseParty):
 
     """
 
-    def __init__(self, client, path, identifier=None):
+    def __init__(
+        self, client: KazooClient, path: str, identifier: Optional[str] = None
+    ):
         BaseParty.__init__(self, client, path, identifier=identifier)
         self.node = "-".join([uuid.uuid4().hex, self.data.decode("utf-8")])
-        self.create_path = self.path + "/" + self.node
+        self._create_path = self.path + "/" + self.node
+
+    @property
+    def create_path(self) -> str:
+        return self._create_path
 
-    def __iter__(self):
+    def __iter__(self) -> Generator[str, None, None]:
         """Get a list of participating clients' identifiers"""
         self._ensure_parent()
         children = self._get_children()

kazoo/retry.py

@@ -3,6 +3,7 @@
 import logging
 import random
 import time
+from typing import TYPE_CHECKING
 
 from kazoo.exceptions import (
     ConnectionClosedError,
@@ -13,6 +14,14 @@
 )
 
 
+if TYPE_CHECKING:
+    from typing import Callable, TypeVar
+    from typing_extensions import ParamSpec
+
+    P = ParamSpec("P")
+    R = TypeVar("R")
+
+
 log = logging.getLogger(__name__)
 
 
@@ -111,7 +120,9 @@ def copy(self):
         obj.retry_exceptions = self.retry_exceptions
         return obj
 
-    def __call__(self, func, *args, **kwargs):
+    def __call__(
+        self, func: Callable[P, R], *args: P.args, **kwargs: P.kwargs
+    ) -> R:
         """Call a function with arguments until it completes without
         throwing a Kazoo exception
 

pyproject.toml

@@ -104,7 +104,6 @@ module = [
     'kazoo.recipe.cache',
     'kazoo.recipe.lock',
     'kazoo.recipe.partitioner',
-    'kazoo.recipe.party',
     'kazoo.recipe.queue',
     'kazoo.recipe.watchers',
     'kazoo.retry',