From 10cd77943bdf9bee4868ee66397a4ebf0443bd28 Mon Sep 17 00:00:00 2001
From: Alexey Kotenko <alexey@kotenko.me>
Date: Wed, 24 Jul 2024 12:49:38 +0200
Subject: [PATCH] Improve adapters/django typings

---
 theine/adapters/django.py | 39 ++++++++++++++++++++-------------------
 1 file changed, 20 insertions(+), 19 deletions(-)

diff --git a/theine/adapters/django.py b/theine/adapters/django.py
index 902b4ed..65437e1 100644
--- a/theine/adapters/django.py
+++ b/theine/adapters/django.py
@@ -1,45 +1,46 @@
 from datetime import timedelta
 from threading import Lock
-from typing import Optional, cast
+from typing import Any, Callable, Optional, Union, cast, Mapping
 
 from django.core.cache.backends.base import DEFAULT_TIMEOUT, BaseCache
 
 from theine import Cache as Theine
 from theine.theine import sentinel
 
+KEY_TYPE = Union[str, Callable[..., str]]
+VALUE_TYPE = Any
+VERSION_TYPE = Optional[int]
+
 
 class Cache(BaseCache):
-    def __init__(self, name, params):
+    def __init__(self, name: str, params: dict[str, Any]):
         super().__init__(params)
         options = params.get("OPTIONS", {})
         policy = options.get("POLICY", "tlfu")
         self.cache = Theine(policy, self._max_entries)
 
-    def _timeout_seconds(self, timeout) -> Optional[float]:
-        if timeout == DEFAULT_TIMEOUT:
-            return self.default_timeout
-        return timeout
-
-    def add(self, key, value, timeout=DEFAULT_TIMEOUT, version=None):
+    def add(self, key: KEY_TYPE, value: VALUE_TYPE, timeout: Optional[float] = DEFAULT_TIMEOUT,
+            version: VERSION_TYPE = None) -> bool:
         data = self.get(key, sentinel, version)
         if data is not sentinel:
             return False
         key = self.make_key(key, version)
+        backend_timeout = self.get_backend_timeout(timeout)
         self.cache.set(
             key,
             value,
-            timedelta(seconds=cast(float, timeout))
-            if timeout is not DEFAULT_TIMEOUT
-            else None,
+            timedelta(seconds=backend_timeout) if backend_timeout else None,
         )
         return True
 
-    def get(self, key, default=None, version=None):
+    def get(self, key: KEY_TYPE, default: Optional[VALUE_TYPE] = None, version: VERSION_TYPE = None) -> Optional[
+        VALUE_TYPE]:
         key = self.make_key(key, version)
         return self.cache.get(key, default)
 
-    def set(self, key, value, timeout=DEFAULT_TIMEOUT, version=None):
-        to = self._timeout_seconds(timeout)
+    def set(self, key: KEY_TYPE, value: VALUE_TYPE, timeout: Optional[float] = DEFAULT_TIMEOUT,
+            version: VERSION_TYPE = None) -> None:
+        to = self.get_backend_timeout(timeout)
         if to is not None and to <= 0:
             self.delete(key)
             return
@@ -50,7 +51,7 @@ def set(self, key, value, timeout=DEFAULT_TIMEOUT, version=None):
             timedelta(seconds=to) if to is not None else None,
         )
 
-    def touch(self, key, timeout=DEFAULT_TIMEOUT, version=None):
+    def touch(self, key: KEY_TYPE, timeout: Optional[float] = DEFAULT_TIMEOUT, version: VERSION_TYPE = None) -> bool:
         data = self.get(key, sentinel, version)
         if data is sentinel:
             return False
@@ -58,18 +59,18 @@ def touch(self, key, timeout=DEFAULT_TIMEOUT, version=None):
         if (
             timeout is not DEFAULT_TIMEOUT
             and timeout is not None
-            and cast(float, timeout) <= 0
+            and timeout <= 0
         ):
             self.cache.delete(nkey)
             return True
-        to = self._timeout_seconds(timeout)
+        to = self.get_backend_timeout(timeout)
         with Lock():
             self.cache._access(nkey, timedelta(seconds=to) if to is not None else None)
         return True
 
-    def delete(self, key, version=None):
+    def delete(self, key: KEY_TYPE, version: VERSION_TYPE = None) -> bool:
         key = self.make_key(key, version)
         return self.cache.delete(key)
 
-    def clear(self):
+    def clear(self) -> None:
         self.cache.clear()