restructure of type hint system of pools for better mypy compatibility

Author: amirreza8002

django_valkey/async_cache/pool.py

@@ -12,7 +12,7 @@
 from django_valkey.typings import AsyncDefaultParserT
 
 
-class AsyncConnectionFactory(BaseConnectionFactory[AValkey, ConnectionPool]):
+class AsyncConnectionFactory(BaseConnectionFactory[AValkey]):
     path_pool_cls = "valkey.asyncio.connection.ConnectionPool"
     path_base_cls = "valkey.asyncio.client.Valkey"
 

django_valkey/base_pool.py

@@ -3,25 +3,24 @@
 from django.core.exceptions import ImproperlyConfigured
 from django.utils.module_loading import import_string
 
-Pool = TypeVar("Pool")
 Base = TypeVar("Base")
 
 
-class BaseConnectionFactory(Generic[Base, Pool]):
+class BaseConnectionFactory(Generic[Base]):
     # Store connection pool by cache backend options.
     #
     # _pools is a process-global, as otherwise _pools is cleared every time
     # ConnectionFactory is instantiated, as Django creates new cache client
     # (DefaultClient) instance for every request.
 
-    _pools: dict[str, Pool | Any] = {}
+    _pools: dict[str, type] = {}  # dict[str, Pool]
 
     def __init__(self, options: dict) -> None:
-        pool_cls_path: str = options.get("CONNECTION_POOL_CLASS", self.path_pool_cls)
-        self.pool_cls: type[Pool] = import_string(pool_cls_path)
+        pool_cls_path: str = options.get("CONNECTION_POOL_CLASS", self.path_pool_cls)  # type: ignore[attr-defined]
+        self.pool_cls = import_string(pool_cls_path)  # type[Pool]
         self.pool_cls_kwargs = options.get("CONNECTION_POOL_KWARGS", {})
 
-        base_client_cls_path: str = options.get("BASE_CLIENT_CLASS", self.path_base_cls)
+        base_client_cls_path: str = options.get("BASE_CLIENT_CLASS", self.path_base_cls)  # type: ignore[attr-defined]
         self.base_client_cls: type[Base] = import_string(base_client_cls_path)
         self.base_client_cls_kwargs = options.get("BASE_CLIENT_KWARGS", {})
 
@@ -42,7 +41,7 @@ def make_connection_params(self, url: str | None) -> dict:
         # TODO: do we need to check for existence?
         if socket_timeout:
             if not isinstance(socket_timeout, (int, float)):
-                error_message = "Socket timeout should be float or integer"
+                error_message = "Socket timeout should be float or integer"  # type: ignore[unreachable]
                 raise ImproperlyConfigured(error_message)
             kwargs["socket_timeout"] = socket_timeout
 
@@ -51,7 +50,7 @@ def make_connection_params(self, url: str | None) -> dict:
         )
         if socket_connect_timeout:
             if not isinstance(socket_connect_timeout, (int, float)):
-                error_message = "Socket connect timeout should be float or integer"
+                error_message = "Socket connect timeout should be float or integer"  # type: ignore[unreachable]
                 raise ImproperlyConfigured(error_message)
             kwargs["socket_connect_timeout"] = socket_connect_timeout
 
@@ -61,31 +60,34 @@ def make_connection_params(self, url: str | None) -> dict:
 
         return kwargs
 
-    def get_connection_pool(self, params: dict) -> Pool:
+    def get_connection_pool(self, params: dict):
         """
         Given a connection parameters, return a new
         connection pool for them.
 
         Overwrite this method if you want a custom
         behavior on creating connection pool.
+
+        :returns: the connection pool
         """
-        cp_params = params
-        cp_params.update(self.pool_cls_kwargs)
-        pool = self.pool_cls.from_url(**cp_params)
+        params.update(self.pool_cls_kwargs)
+        pool = self.pool_cls.from_url(**params)
 
         if pool.connection_kwargs.get("password", None) is None:
             pool.connection_kwargs["password"] = params.get("password")
             pool.reset()
 
         return pool
 
-    def get_or_create_connection_pool(self, params: dict) -> Pool:
+    def get_or_create_connection_pool(self, params: dict):
         """
         Given a connection parameters and return a new
         or cached connection pool for them.
 
         Reimplement this method if you want distinct
         connection pool instance caching behavior.
+
+        :returns: the connection pool
         """
         key: str = params["url"]
         if key not in self._pools:

django_valkey/cluster_cache/pool.py

@@ -7,7 +7,7 @@
 from django_valkey.base_pool import BaseConnectionFactory
 
 
-class ClusterConnectionFactory(BaseConnectionFactory[ValkeyCluster, ConnectionPool]):
+class ClusterConnectionFactory(BaseConnectionFactory[ValkeyCluster]):
     path_pool_cls = "valkey.connection.ConnectionPool"
     path_base_cls = "valkey.cluster.ValkeyCluster"
 

django_valkey/pool.py

@@ -1,23 +1,27 @@
+from typing import TYPE_CHECKING
 from urllib.parse import parse_qs, urlparse
 
 from django.conf import settings
 from django.core.exceptions import ImproperlyConfigured
 from django.utils.module_loading import import_string
 from valkey import Valkey
-from valkey.connection import ConnectionPool, DefaultParser
+from valkey.connection import DefaultParser
 from valkey.sentinel import Sentinel
 from valkey._parsers.url_parser import to_bool
 
 from django_valkey.base_pool import BaseConnectionFactory
 from django_valkey.typings import DefaultParserT
 
-from django_valkey.async_cache.pool import (
-    AsyncConnectionFactory,
-    AsyncSentinelConnectionFactory,
-)
+if TYPE_CHECKING:
+    from django_valkey.async_cache.pool import (
+        AsyncConnectionFactory,
+        AsyncSentinelConnectionFactory,
+    )
+    from valkey.connection import ConnectionPool
+    from valkey.sentinel import SentinelConnectionPool
 
 
-class ConnectionFactory(BaseConnectionFactory[Valkey, ConnectionPool]):
+class ConnectionFactory(BaseConnectionFactory[Valkey]):
     path_pool_cls = "valkey.connection.ConnectionPool"
     path_base_cls = "valkey.client.Valkey"
 
@@ -40,11 +44,13 @@ def connect(self, url: str) -> Valkey:
         return self.get_connection(params)
 
     def get_connection(self, params: dict) -> Valkey:
-        pool = self.get_or_create_connection_pool(params)
+        pool: ConnectionPool = self.get_or_create_connection_pool(params)
         return self.base_client_cls(connection_pool=pool, **self.base_client_cls_kwargs)
 
 
 class SentinelConnectionFactory(ConnectionFactory):
+    path_pool_cls = "valkey.sentinel.SentinelConnectionPool"
+
     def __init__(self, options: dict):
         # allow overriding the default SentinelConnectionPool class
         options.setdefault(
@@ -68,7 +74,7 @@ def __init__(self, options: dict):
             **connection_kwargs,
         )
 
-    def get_connection_pool(self, params: dict) -> ConnectionPool:
+    def get_connection_pool(self, params: dict) -> "SentinelConnectionPool":
         """
         Given a connection parameters, return a new sentinel connection pool
         for them.
@@ -79,7 +85,7 @@ def get_connection_pool(self, params: dict) -> ConnectionPool:
         # SentinelConnectionPool constructor since will be called by from_url
         cp_params = params
         cp_params.update(service_name=url.hostname, sentinel_manager=self._sentinel)
-        pool = super().get_connection_pool(cp_params)
+        pool: SentinelConnectionPool = super().get_connection_pool(cp_params)
 
         # convert "is_master" to a boolean if set on the URL, otherwise if not
         # provided it defaults to True.
@@ -92,12 +98,7 @@ def get_connection_pool(self, params: dict) -> ConnectionPool:
 
 def get_connection_factory(
     path: str | None = None, options: dict | None = None
-) -> (
-    ConnectionFactory
-    | SentinelConnectionFactory
-    | AsyncConnectionFactory
-    | AsyncSentinelConnectionFactory
-):
+) -> "ConnectionFactory | SentinelConnectionFactory | AsyncConnectionFactory | AsyncSentinelConnectionFactory":
 
     path = getattr(settings, "DJANGO_VALKEY_CONNECTION_FACTORY", path)
     if options: