docs: add types to recipe.counter

Author: a-ungurianu

kazoo/client.py

@@ -1122,7 +1122,7 @@ def _create_async_inner(
             raise async_result.exception
         return async_result
 
-    def ensure_path(self, path, acl=None):
+    def ensure_path(self, path: str, acl: Optional[List[ACL]] = None) -> bool:
         """Recursively create a path if it doesn't exist.
 
         :param path: Path of node.
@@ -1211,7 +1211,9 @@ def exists_async(self, path, watch=None):
         )
         return async_result
 
-    def get(self, path, watch=None):
+    def get(
+        self, path: str, watch: Optional[WatchListener] = None
+    ) -> Tuple[bytes, ZnodeStat]:
         """Get the value of a node.
 
         If a watch is provided, it will be left on the node with the
@@ -1400,7 +1402,7 @@ def set_acls_async(self, path, acls, version=-1):
         )
         return async_result
 
-    def set(self, path, value, version=-1):
+    def set(self, path: str, value: bytes, version: int = -1) -> ZnodeStat:
         """Set the value of a node.
 
         If the version of the node being updated is newer than the

kazoo/recipe/counter.py

@@ -4,9 +4,20 @@
 :Status: Unknown
 
 """
+from __future__ import annotations
+
+import struct
+from typing import cast, TYPE_CHECKING
+
 from kazoo.exceptions import BadVersionError
 from kazoo.retry import ForceRetryError
-import struct
+
+if TYPE_CHECKING:
+    from typing import Optional, Tuple, Type, Union
+
+    from kazoo.client import KazooClient
+
+    CountT = Union[int, float]
 
 
 class Counter(object):
@@ -58,7 +69,13 @@ class Counter(object):
 
     """
 
-    def __init__(self, client, path, default=0, support_curator=False):
+    def __init__(
+        self,
+        client: KazooClient,
+        path: str,
+        default: CountT = 0,
+        support_curator: bool = False,
+    ):
         """Create a Kazoo Counter
 
         :param client: A :class:`~kazoo.client.KazooClient` instance.
@@ -70,46 +87,50 @@ def __init__(self, client, path, default=0, support_curator=False):
         """
         self.client = client
         self.path = path
-        self.default = default
-        self.default_type = type(default)
+        self.default: CountT = default
+        self.default_type: Type[CountT] = type(default)
         self.support_curator = support_curator
         self._ensured_path = False
-        self.pre_value = None
-        self.post_value = None
+        self.pre_value: Optional[CountT] = None
+        self.post_value: Optional[CountT] = None
         if self.support_curator and not isinstance(self.default, int):
             raise TypeError(
                 "when support_curator is enabled the default "
                 "type must be an int"
             )
 
-    def _ensure_node(self):
+    def _ensure_node(self) -> None:
         if not self._ensured_path:
             # make sure our node exists
             self.client.ensure_path(self.path)
             self._ensured_path = True
 
-    def _value(self):
+    def _value(self) -> Tuple[CountT, int]:
         self._ensure_node()
         old, stat = self.client.get(self.path)
         if self.support_curator:
-            old = struct.unpack(">i", old)[0] if old != b"" else self.default
+            parsed_old: Union[int, float, str] = (
+                cast(int, struct.unpack(">i", old)[0])
+                if old != b""
+                else self.default
+            )
         else:
-            old = old.decode("ascii") if old != b"" else self.default
+            parsed_old = old.decode("ascii") if old != b"" else self.default
         version = stat.version
-        data = self.default_type(old)
+        data = self.default_type(parsed_old)
         return data, version
 
     @property
-    def value(self):
+    def value(self) -> CountT:
         return self._value()[0]
 
-    def _change(self, value):
+    def _change(self, value: CountT) -> "Counter":
         if not isinstance(value, self.default_type):
             raise TypeError("invalid type for value change")
         self.client.retry(self._inner_change, value)
         return self
 
-    def _inner_change(self, value):
+    def _inner_change(self, value: CountT) -> None:
         self.pre_value, version = self._value()
         post_value = self.pre_value + value
         if self.support_curator:
@@ -123,10 +144,10 @@ def _inner_change(self, value):
             raise ForceRetryError()
         self.post_value = post_value
 
-    def __add__(self, value):
+    def __add__(self, value: CountT) -> "Counter":
         """Add value to counter."""
         return self._change(value)
 
-    def __sub__(self, value):
+    def __sub__(self, value: CountT) -> "Counter":
         """Subtract value from counter."""
         return self._change(-value)

pyproject.toml

@@ -103,7 +103,6 @@ module = [
     'kazoo.protocol.states',
     'kazoo.recipe.barrier',
     'kazoo.recipe.cache',
-    'kazoo.recipe.counter',
     'kazoo.recipe.lock',
     'kazoo.recipe.partitioner',
     'kazoo.recipe.party',