docs: add types to recipe.lease

Author: a-ungurianu

kazoo/recipe/lease.py

@@ -5,12 +5,27 @@
 :Status: Beta
 
 """
+from __future__ import annotations
+
 import datetime
 import json
 import socket
+from typing import cast, TYPE_CHECKING
 
 from kazoo.exceptions import CancelledError
 
+from typing_extensions import TypedDict
+
+if TYPE_CHECKING:
+    from kazoo.client import KazooClient
+    from typing import Callable, Optional
+
+
+class LeaseData(TypedDict):
+    version: int
+    holder: str
+    end: str
+
 
 class NonBlockingLease(object):
     """Exclusive lease that does not block.
@@ -48,11 +63,11 @@ class NonBlockingLease(object):
 
     def __init__(
         self,
-        client,
-        path,
-        duration,
-        identifier=None,
-        utcnow=datetime.datetime.utcnow,
+        client: KazooClient,
+        path: str,
+        duration: datetime.timedelta,
+        identifier: Optional[str] = None,
+        utcnow: Callable[[], datetime.datetime] = datetime.datetime.utcnow,
     ):
         """Create a non-blocking lease.
 
@@ -71,7 +86,14 @@ def __init__(
         self.obtained = False
         self._attempt_obtaining(client, path, duration, ident, utcnow)
 
-    def _attempt_obtaining(self, client, path, duration, ident, utcnow):
+    def _attempt_obtaining(
+        self,
+        client: KazooClient,
+        path: str,
+        duration: datetime.timedelta,
+        ident: str,
+        utcnow: Callable[[], datetime.datetime],
+    ) -> None:
         client.ensure_path(path)
         holder_path = path + "/lease_holder"
         lock = client.Lock(path, ident)
@@ -92,7 +114,7 @@ def _attempt_obtaining(self, client, path, duration, ident, utcnow):
                         return
                     client.delete(holder_path)
                 end_lease = (now + duration).strftime(self._date_format)
-                new_data = {
+                new_data: LeaseData = {
                     "version": self._version,
                     "holder": ident,
                     "end": end_lease,
@@ -103,18 +125,18 @@ def _attempt_obtaining(self, client, path, duration, ident, utcnow):
         except CancelledError:
             pass
 
-    def _encode(self, data_dict):
+    def _encode(self, data_dict: LeaseData) -> bytes:
         return json.dumps(data_dict).encode(self._byte_encoding)
 
-    def _decode(self, raw):
-        return json.loads(raw.decode(self._byte_encoding))
+    def _decode(self, raw: bytes) -> LeaseData:
+        return cast(LeaseData, json.loads(raw.decode(self._byte_encoding)))
 
     # Python 2.x
-    def __nonzero__(self):
+    def __nonzero__(self) -> bool:
         return self.obtained
 
     # Python 3.x
-    def __bool__(self):
+    def __bool__(self) -> bool:
         return self.obtained
 
 
@@ -140,12 +162,12 @@ class MultiNonBlockingLease(object):
 
     def __init__(
         self,
-        client,
-        count,
-        path,
-        duration,
-        identifier=None,
-        utcnow=datetime.datetime.utcnow,
+        client: KazooClient,
+        count: int,
+        path: str,
+        duration: datetime.timedelta,
+        identifier: Optional[str] = None,
+        utcnow: Callable[[], datetime.datetime] = datetime.datetime.utcnow,
     ):
         self.obtained = False
         for num in range(count):
@@ -161,9 +183,9 @@ def __init__(
                 break
 
     # Python 2.x
-    def __nonzero__(self):
+    def __nonzero__(self) -> bool:
         return self.obtained
 
     # Python 3.x
-    def __bool__(self):
+    def __bool__(self) -> bool:
         return self.obtained

pyproject.toml

@@ -104,7 +104,6 @@ module = [
     'kazoo.recipe.barrier',
     'kazoo.recipe.cache',
     'kazoo.recipe.counter',
-    'kazoo.recipe.lease',
     'kazoo.recipe.lock',
     'kazoo.recipe.partitioner',
     'kazoo.recipe.party',

setup.cfg

@@ -39,6 +39,8 @@ project_urls =
 zip_safe = false
 include_package_data = true
 packages = find:
+install_requires =
+    typing-extensions
 
 [aliases]
 release = sdist bdist_wheel
@@ -83,4 +85,3 @@ alldeps =
     %(sasl)s
     %(docs)s
     %(typing)s
-