docs: add types to recipe.lock

Author: a-ungurianu

kazoo/client.py

@@ -917,7 +917,7 @@ def create(
         sequence=False,
         makepath=False,
         include_data=False,
-    ):
+    ) -> str:
         """Create a node with the given value as its data. Optionally
         set an ACL on the node.
 

kazoo/protocol/states.py

@@ -1,8 +1,9 @@
 """Kazoo State and Event objects"""
 from collections import namedtuple
+from enum import Enum
 
 
-class KazooState(object):
+class KazooState(str, Enum):
     """High level connection state values
 
     States inspired by Netflix Curator.

kazoo/recipe/lock.py

@@ -15,40 +15,37 @@
 
 """
 import re
-from time import monotonic as now
 import uuid
+from time import monotonic as now
+from typing import TYPE_CHECKING, Any, ContextManager, List, Optional, Sequence, Type, Union, cast
+
+from kazoo.exceptions import (CancelledError, KazooException, LockTimeout,
+                              NoNodeError)
+from kazoo.protocol.states import KazooState, WatchedEvent
+from kazoo.retry import ForceRetryError, KazooRetry, RetryFailedError
 
-from kazoo.exceptions import (
-    CancelledError,
-    KazooException,
-    LockTimeout,
-    NoNodeError,
-)
-from kazoo.protocol.states import KazooState
-from kazoo.retry import (
-    ForceRetryError,
-    KazooRetry,
-    RetryFailedError,
-)
+if TYPE_CHECKING:
+    from kazoo.client import KazooClient
+    from typing_extensions import Literal
 
 
 class _Watch(object):
-    def __init__(self, duration=None):
+    def __init__(self, duration: Optional[float] = None):
         self.duration = duration
-        self.started_at = None
+        self.started_at: Optional[float] = None
 
-    def start(self):
+    def start(self) -> None:
         self.started_at = now()
 
-    def leftover(self):
+    def leftover(self) -> Optional[float]:
         if self.duration is None:
             return None
         else:
-            elapsed = now() - self.started_at
+            elapsed = now() - cast(float, self.started_at)
             return max(0, self.duration - elapsed)
 
 
-class Lock(object):
+class Lock(ContextManager[None]):
     """Kazoo Lock
 
     Example usage with a :class:`~kazoo.client.KazooClient` instance:
@@ -77,7 +74,13 @@ class Lock(object):
     # sequence number. Involved in read/write locks.
     _EXCLUDE_NAMES = ["__lock__"]
 
-    def __init__(self, client, path, identifier=None, extra_lock_patterns=()):
+    def __init__(
+        self,
+        client: KazooClient,
+        path: str,
+        identifier: Optional[str] = None,
+        extra_lock_patterns: Sequence[str] = (),
+    ):
         """Create a Kazoo lock.
 
         :param client: A :class:`~kazoo.client.KazooClient` instance.
@@ -109,7 +112,7 @@ def __init__(self, client, path, identifier=None, extra_lock_patterns=()):
         # some data is written to the node. this can be queried via
         # contenders() to see who is contending for the lock
         self.data = str(identifier or "").encode("utf-8")
-        self.node = None
+        self.node: Optional[str] = None
 
         self.wake_event = client.handler.event_object()
 
@@ -125,20 +128,25 @@ def __init__(self, client, path, identifier=None, extra_lock_patterns=()):
         self.assured_path = False
         self.cancelled = False
         self._retry = KazooRetry(
-            max_tries=None, sleep_func=client.handler.sleep_func
+            max_tries=-1, sleep_func=client.handler.sleep_func
         )
         self._acquire_method_lock = client.handler.lock_object()
 
-    def _ensure_path(self):
+    def _ensure_path(self) -> None:
         self.client.ensure_path(self.path)
         self.assured_path = True
 
-    def cancel(self):
+    def cancel(self) -> None:
         """Cancel a pending lock acquire."""
         self.cancelled = True
         self.wake_event.set()
 
-    def acquire(self, blocking=True, timeout=None, ephemeral=True):
+    def acquire(
+        self,
+        blocking: bool = True,
+        timeout: Optional[float] = None,
+        ephemeral: bool = True,
+    ) -> bool:
         """
         Acquire the lock. By defaults blocks and waits forever.
 
@@ -204,11 +212,13 @@ def acquire(self, blocking=True, timeout=None, ephemeral=True):
         finally:
             self._acquire_method_lock.release()
 
-    def _watch_session(self, state):
+    def _watch_session(self, state: KazooState) -> bool:
         self.wake_event.set()
         return True
 
-    def _inner_acquire(self, blocking, timeout, ephemeral=True):
+    def _inner_acquire(
+        self, blocking: bool, timeout: Optional[float], ephemeral: bool=True
+    ) -> bool:
 
         # wait until it's our chance to get it..
         if self.is_acquired:
@@ -266,10 +276,10 @@ def _inner_acquire(self, blocking, timeout, ephemeral=True):
             finally:
                 self.client.remove_listener(self._watch_session)
 
-    def _watch_predecessor(self, event):
+    def _watch_predecessor(self, event: WatchedEvent) -> None:
         self.wake_event.set()
 
-    def _get_predecessor(self, node):
+    def _get_predecessor(self, node: str) -> Optional[str]:
         """returns `node`'s predecessor or None
 
         Note: This handle the case where the current lock is not a contender
@@ -278,7 +288,7 @@ def _get_predecessor(self, node):
         """
         node_sequence = node[len(self.prefix) :]
         children = self.client.get_children(self.path)
-        found_self = False
+        found_self: Union[Literal[False],re.Match[bytes]] = False
         # Filter out the contenders using the computed regex
         contender_matches = []
         for child in children:
@@ -293,7 +303,7 @@ def _get_predecessor(self, node):
             if child == node:
                 # Remember the node's match object so we can short circuit
                 # below.
-                found_self = match
+                found_self = cast(re.Match[bytes], match)
 
         if found_self is False:  # pragma: nocover
             # somehow we aren't in the childrens -- probably we are
@@ -309,42 +319,42 @@ def _get_predecessor(self, node):
         sorted_matches = sorted(contender_matches, key=lambda m: m.groups())
         return sorted_matches[-1].string
 
-    def _find_node(self):
+    def _find_node(self) -> Optional[str]:
         children = self.client.get_children(self.path)
         for child in children:
             if child.startswith(self.prefix):
                 return child
         return None
 
-    def _delete_node(self, node):
+    def _delete_node(self, node: str) -> None:
         self.client.delete(self.path + "/" + node)
 
-    def _best_effort_cleanup(self):
+    def _best_effort_cleanup(self) -> None:
         try:
             node = self.node or self._find_node()
             if node:
                 self._delete_node(node)
         except KazooException:  # pragma: nocover
             pass
 
-    def release(self):
+    def release(self) -> bool:
         """Release the lock immediately."""
         return self.client.retry(self._inner_release)
 
-    def _inner_release(self):
+    def _inner_release(self) -> bool:
         if not self.is_acquired:
             return False
 
         try:
-            self._delete_node(self.node)
+            self._delete_node(cast(str, self.node))
         except NoNodeError:  # pragma: nocover
             pass
 
         self.is_acquired = False
         self.node = None
         return True
 
-    def contenders(self):
+    def contenders(self) -> List[str]:
         """Return an ordered list of the current contenders for the
         lock.
 
@@ -380,7 +390,7 @@ def contenders(self):
             for match in sorted(contender_matches, key=lambda m: m.groups())
         ]
         # Retrieve all the contender nodes data (preserving order).
-        contenders = []
+        contenders: List[str] = []
         for node in contender_nodes:
             try:
                 data, stat = self.client.get(self.path + "/" + node)
@@ -391,10 +401,10 @@ def contenders(self):
 
         return contenders
 
-    def __enter__(self):
+    def __enter__(self) -> None:
         self.acquire()
 
-    def __exit__(self, exc_type, exc_value, traceback):
+    def __exit__(self, exc_type: Optional[Type[BaseException]], exc_value: Optional[BaseException], traceback: Any) -> None:
         self.release()
 
 
@@ -458,7 +468,7 @@ class ReadLock(Lock):
     _EXCLUDE_NAMES = ["__lock__"]
 
 
-class Semaphore(object):
+class Semaphore(ContextManager[None]):
     """A Zookeeper-based Semaphore
 
     This synchronization primitive operates in the same manner as the
@@ -493,7 +503,13 @@ class Semaphore(object):
 
     """
 
-    def __init__(self, client, path, identifier=None, max_leases=1):
+    def __init__(
+        self,
+        client: KazooClient,
+        path: str,
+        identifier: Optional[str] = None,
+        max_leases: int = 1,
+    ):
         """Create a Kazoo Lock
 
         :param client: A :class:`~kazoo.client.KazooClient` instance.
@@ -529,7 +545,7 @@ def __init__(self, client, path, identifier=None, max_leases=1):
         self.cancelled = False
         self._session_expired = False
 
-    def _ensure_path(self):
+    def _ensure_path(self) -> None:
         result = self.client.ensure_path(self.path)
         self.assured_path = True
         if result is True:
@@ -550,12 +566,14 @@ def _ensure_path(self):
         else:
             self.client.set(self.path, str(self.max_leases).encode("utf-8"))
 
-    def cancel(self):
+    def cancel(self) -> None:
         """Cancel a pending semaphore acquire."""
         self.cancelled = True
         self.wake_event.set()
 
-    def acquire(self, blocking=True, timeout=None):
+    def acquire(
+        self, blocking: bool = True, timeout: Optional[float] = None
+    ) -> bool:
         """Acquire the semaphore. By defaults blocks and waits forever.
 
         :param blocking: Block until semaphore is obtained or
@@ -593,7 +611,7 @@ def acquire(self, blocking=True, timeout=None):
 
         return self.is_acquired
 
-    def _inner_acquire(self, blocking, timeout=None):
+    def _inner_acquire(self, blocking: bool, timeout: Optional[float] = None) -> bool:
         """Inner loop that runs from the top anytime a command hits a
         retryable Zookeeper exception."""
         self._session_expired = False
@@ -634,10 +652,10 @@ def _inner_acquire(self, blocking, timeout=None):
         finally:
             lock.release()
 
-    def _watch_lease_change(self, event):
+    def _watch_lease_change(self, event: WatchedEvent) -> None:
         self.wake_event.set()
 
-    def _get_lease(self, data=None):
+    def _get_lease(self) -> bool:
         # Make sure the session is still valid
         if self._session_expired:
             raise ForceRetryError("Retry on session loss at top")
@@ -666,25 +684,26 @@ def _get_lease(self, data=None):
         # Return current state
         return self.is_acquired
 
-    def _watch_session(self, state):
+    def _watch_session(self, state: KazooState) -> Optional[Literal[True]]:
         if state == KazooState.LOST:
             self._session_expired = True
             self.wake_event.set()
 
             # Return true to de-register
             return True
+        return None
 
-    def _best_effort_cleanup(self):
+    def _best_effort_cleanup(self) -> None:
         try:
             self.client.delete(self.create_path)
         except KazooException:  # pragma: nocover
             pass
 
-    def release(self):
+    def release(self) -> bool:
         """Release the lease immediately."""
         return self.client.retry(self._inner_release)
 
-    def _inner_release(self):
+    def _inner_release(self) -> bool:
         if not self.is_acquired:
             return False
         try:
@@ -694,7 +713,7 @@ def _inner_release(self):
         self.is_acquired = False
         return True
 
-    def lease_holders(self):
+    def lease_holders(self) -> List[str]:
         """Return an unordered list of the current lease holders.
 
         .. note::
@@ -708,7 +727,7 @@ def lease_holders(self):
 
         children = self.client.get_children(self.path)
 
-        lease_holders = []
+        lease_holders: List[str] = []
         for child in children:
             try:
                 data, stat = self.client.get(self.path + "/" + child)
@@ -717,8 +736,8 @@ def lease_holders(self):
                 pass
         return lease_holders
 
-    def __enter__(self):
+    def __enter__(self) -> None:
         self.acquire()
 
-    def __exit__(self, exc_type, exc_value, traceback):
+    def __exit__(self, exc_type: Optional[Type[BaseException]], exc_value: Optional[BaseException], traceback: Any) -> None:
         self.release()

pyproject.toml

@@ -29,7 +29,8 @@ follow_imports = "silent"
 ignore_missing_imports = true
 disable_error_code = ["no-untyped-call", "no-untyped-def"]
 files = [
-    "kazoo/recipe/election.py"
+    "kazoo/recipe/election.py",
+    "kazoo/recipe/lock.py"
 ]
 
 [[tool.mypy.overrides]]