docs: add types to recipe.lock

Author: a-ungurianu

kazoo/client.py

@@ -922,7 +922,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

@@ -14,14 +14,17 @@
 and/or the lease has been lost.
 
 """
+from __future__ import annotations
+
 import re
 import sys
+from typing import TYPE_CHECKING, cast, ContextManager
+import uuid
 
 try:
     from time import monotonic as now
 except ImportError:
     from time import time as now
-import uuid
 
 import six
 
@@ -38,24 +41,37 @@
     RetryFailedError,
 )
 
+if TYPE_CHECKING:
+    from kazoo.client import KazooClient
+    from kazoo.protocol.states import WatchedEvent
+    from typing import (
+        Any,
+        List,
+        Optional,
+        Sequence,
+        Type,
+        Union,
+    )
+    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:
@@ -84,7 +100,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.
@@ -116,7 +138,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()
 
@@ -132,20 +154,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.
 
@@ -212,11 +239,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:
@@ -274,10 +303,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
@@ -286,7 +315,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:
@@ -301,7 +330,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
@@ -317,42 +346,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.
 
@@ -388,7 +417,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)
@@ -399,10 +428,15 @@ 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()
 
 
@@ -466,7 +500,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
@@ -501,7 +535,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.
@@ -537,7 +577,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:
@@ -558,12 +598,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
@@ -601,7 +643,9 @@ 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
@@ -642,10 +686,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")
@@ -674,25 +718,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:
@@ -702,7 +747,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::
@@ -716,7 +761,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)
@@ -725,8 +770,13 @@ 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()