docs: add types to recipe.barrier

Author: a-ungurianu

kazoo/client.py

@@ -5,6 +5,7 @@
 import logging
 from os.path import split
 import re
+from typing import TYPE_CHECKING
 import warnings
 
 from kazoo.exceptions import (
@@ -63,6 +64,20 @@
 from kazoo.recipe.queue import Queue, LockingQueue
 from kazoo.recipe.watchers import ChildrenWatch, DataWatch
 
+if TYPE_CHECKING:
+    from typing import (
+        List,
+        Optional,
+        Sequence,
+        Tuple,
+        Union,
+        Callable,
+        Literal,
+        overload,
+    )
+    from kazoo.protocol.states import ZnodeStat
+
+    WatchListener = Callable[[WatchedEvent], None]
 
 CLOSED_STATES = (
     KeeperState.EXPIRED_SESSION,
@@ -910,14 +925,14 @@ def sync(self, path):
 
     def create(
         self,
-        path,
-        value=b"",
-        acl=None,
-        ephemeral=False,
-        sequence=False,
-        makepath=False,
-        include_data=False,
-    ):
+        path: str,
+        value: bytes = b"",
+        acl: Optional[Sequence[ACL]] = None,
+        ephemeral: bool = False,
+        sequence: bool = False,
+        makepath: bool = False,
+        include_data: bool = False,
+    ) -> Union[str, Tuple[str, ZnodeStat]]:
         """Create a node with the given value as its data. Optionally
         set an ACL on the node.
 
@@ -1171,7 +1186,9 @@ def exists_completion(path, result):
 
         return async_result
 
-    def exists(self, path, watch=None):
+    def exists(
+        self, path: str, watch: Optional[WatchListener] = None
+    ) -> Optional[ZnodeStat]:
         """Check if a node exists.
 
         If a watch is provided, it will be left on the node with the
@@ -1256,7 +1273,53 @@ def get_async(self, path, watch=None):
         )
         return async_result
 
-    def get_children(self, path, watch=None, include_data=False):
+    @overload
+    def get_children(
+        self,
+        path: str,
+    ) -> List[str]:
+        ...
+
+    @overload
+    def get_children(
+        self,
+        path: str,
+        watch: WatchListener,
+    ) -> List[str]:
+        ...
+
+    @overload
+    def get_children(
+        self,
+        path: str,
+        watch: Optional[WatchListener],
+    ) -> List[str]:
+        ...
+
+    @overload
+    def get_children(
+        self,
+        path: str,
+        watch: Optional[WatchListener],
+        include_data: Literal[True],
+    ) -> List[Tuple[str, ZnodeStat]]:
+        ...
+
+    @overload
+    def get_children(
+        self,
+        path: str,
+        watch: Optional[WatchListener] = None,
+        include_data: Literal[False] = False,
+    ) -> List[str]:
+        ...
+
+    def get_children(
+        self,
+        path: str,
+        watch: Optional[WatchListener] = None,
+        include_data: bool = False,
+    ) -> Union[List[Tuple[str, ZnodeStat]], List[str]]:
         """Get a list of child nodes of a path.
 
         If a watch is provided it will be left on the node with the
@@ -1475,7 +1538,9 @@ def transaction(self):
         """
         return TransactionRequest(self)
 
-    def delete(self, path, version=-1, recursive=False):
+    def delete(
+        self, path: str, version: int = -1, recursive: bool = False
+    ) -> Optional[bool]:
         """Delete a node.
 
         The call will succeed if such a node exists, and the given

kazoo/recipe/barrier.py

@@ -4,13 +4,24 @@
 :Status: Unknown
 
 """
+from __future__ import annotations
+
 import os
 import socket
+from threading import Event
+from typing import TYPE_CHECKING, cast
 import uuid
 
 from kazoo.exceptions import KazooException, NoNodeError, NodeExistsError
 from kazoo.protocol.states import EventType
 
+if TYPE_CHECKING:
+    from typing import Optional
+    from typing_extensions import Literal
+
+    from kazoo.client import KazooClient
+    from kazoo.protocol.states import WatchedEvent
+
 
 class Barrier(object):
     """Kazoo Barrier
@@ -27,7 +38,7 @@ class Barrier(object):
 
     """
 
-    def __init__(self, client, path):
+    def __init__(self, client: KazooClient, path: str):
         """Create a Kazoo Barrier
 
         :param client: A :class:`~kazoo.client.KazooClient` instance.
@@ -37,11 +48,11 @@ def __init__(self, client, path):
         self.client = client
         self.path = path
 
-    def create(self):
+    def create(self) -> None:
         """Establish the barrier if it doesn't exist already"""
         self.client.retry(self.client.ensure_path, self.path)
 
-    def remove(self):
+    def remove(self) -> bool:
         """Remove the barrier
 
         :returns: Whether the barrier actually needed to be removed.
@@ -54,17 +65,17 @@ def remove(self):
         except NoNodeError:
             return False
 
-    def wait(self, timeout=None):
+    def wait(self, timeout: Optional[float] = None) -> bool:
         """Wait on the barrier to be cleared
 
         :returns: True if the barrier has been cleared, otherwise
                   False.
         :rtype: bool
 
         """
-        cleared = self.client.handler.event_object()
+        cleared = cast(Event, self.client.handler.event_object())
 
-        def wait_for_clear(event):
+        def wait_for_clear(event: WatchedEvent) -> None:
             if event.type == EventType.DELETED:
                 cleared.set()
 
@@ -93,7 +104,13 @@ class DoubleBarrier(object):
 
     """
 
-    def __init__(self, client, path, num_clients, identifier=None):
+    def __init__(
+        self,
+        client: KazooClient,
+        path: str,
+        num_clients: int,
+        identifier: Optional[str] = None,
+    ):
         """Create a Double Barrier
 
         :param client: A :class:`~kazoo.client.KazooClient` instance.
@@ -118,7 +135,7 @@ def __init__(self, client, path, num_clients, identifier=None):
         self.node_name = uuid.uuid4().hex
         self.create_path = self.path + "/" + self.node_name
 
-    def enter(self):
+    def enter(self) -> None:
         """Enter the barrier, blocks until all nodes have entered"""
         try:
             self.client.retry(self._inner_enter)
@@ -128,7 +145,7 @@ def enter(self):
             self._best_effort_cleanup()
             self.participating = False
 
-    def _inner_enter(self):
+    def _inner_enter(self) -> Literal[True]:
         # make sure our barrier parent node exists
         if not self.assured_path:
             self.client.ensure_path(self.path)
@@ -145,7 +162,7 @@ def _inner_enter(self):
         except NodeExistsError:
             pass
 
-        def created(event):
+        def created(event: WatchedEvent) -> None:
             if event.type == EventType.CREATED:
                 ready.set()
 
@@ -159,7 +176,7 @@ def created(event):
             self.client.ensure_path(self.path + "/ready")
         return True
 
-    def leave(self):
+    def leave(self) -> None:
         """Leave the barrier, blocks until all nodes have left"""
         try:
             self.client.retry(self._inner_leave)
@@ -168,7 +185,7 @@ def leave(self):
             self._best_effort_cleanup()
         self.participating = False
 
-    def _inner_leave(self):
+    def _inner_leave(self) -> Literal[True]:
         # Delete the ready node if its around
         try:
             self.client.delete(self.path + "/ready")
@@ -188,7 +205,7 @@ def _inner_leave(self):
 
             ready = self.client.handler.event_object()
 
-            def deleted(event):
+            def deleted(event: WatchedEvent) -> None:
                 if event.type == EventType.DELETED:
                     ready.set()
 
@@ -214,7 +231,7 @@ def deleted(event):
             # Wait for the lowest to be deleted
             ready.wait()
 
-    def _best_effort_cleanup(self):
+    def _best_effort_cleanup(self) -> None:
         try:
             self.client.retry(self.client.delete, self.create_path)
         except NoNodeError:

pyproject.toml

@@ -101,7 +101,6 @@ module = [
     'kazoo.protocol.paths',
     'kazoo.protocol.serialization',
     'kazoo.protocol.states',
-    'kazoo.recipe.barrier',
     'kazoo.recipe.cache',
     'kazoo.recipe.lock',
     'kazoo.recipe.partitioner',