docs: add types to recipe.barrier

Author: a-ungurianu

kazoo/recipe/barrier.py

@@ -4,13 +4,23 @@
 :Status: Unknown
 
 """
+from __future__ import annotations
+
 import os
 import socket
+from typing import TYPE_CHECKING
 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 +37,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 +47,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,7 +64,7 @@ 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
@@ -64,7 +74,7 @@ def wait(self, timeout=None):
         """
         cleared = 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 +103,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 +134,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 +144,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 +161,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 +175,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 +184,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 +204,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 +230,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

@@ -29,10 +29,13 @@ follow_imports = "silent"
 ignore_missing_imports = true
 disable_error_code = ["no-untyped-call", "no-untyped-def"]
 files = [
+    "kazoo/recipe/barrier.py",
     "kazoo/recipe/election.py",
     "kazoo/recipe/lock.py",
     "kazoo/recipe/lease.py",
-    "kazoo/recipe/counter.py"
+    "kazoo/recipe/party.py",
+    "kazoo/recipe/counter.py",
+    "kazoo/recipe/barrier.py"
 ]
 
 [[tool.mypy.overrides]]