docs: add types to recipe.lease

Author: a-ungurianu

kazoo/recipe/lease.py

@@ -5,12 +5,25 @@
 :Status: Beta
 
 """
+from __future__ import annotations
+
 import datetime
 import json
 import socket
+from typing import cast, TYPE_CHECKING
 
 from kazoo.exceptions import CancelledError
 
+if TYPE_CHECKING:
+    from kazoo.client import KazooClient
+    from typing import Callable, Optional
+    from typing_extensions import TypedDict
+
+    class LeaseData(TypedDict):
+        version: int
+        holder: str
+        end: str
+
 
 class NonBlockingLease(object):
     """Exclusive lease that does not block.
@@ -48,11 +61,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 +84,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 +112,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 +123,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 +160,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 +181,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]]