docs: add types to recipe.lease

Author: a-ungurianu

kazoo/recipe/lease.py

@@ -8,9 +8,20 @@
 import datetime
 import json
 import socket
+from typing import Callable, Optional, cast, TYPE_CHECKING
 
 from kazoo.exceptions import CancelledError
 
+if TYPE_CHECKING: 
+    from typing_extensions import TypedDict
+
+    from kazoo.client import KazooClient
+
+class LeaseData(TypedDict):
+    version: int
+    holder: str
+    end: str
+
 
 class NonBlockingLease(object):
     """Exclusive lease that does not block.
@@ -48,11 +59,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 +82,7 @@ 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 +103,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 +114,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 +151,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 +172,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

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