From c6d34cc24d435cff8f68222cdeab72ec6e500c64 Mon Sep 17 00:00:00 2001
From: Nikita Sobolev <mail@sobolevn.me>
Date: Sat, 20 Aug 2022 20:46:09 +0300
Subject: [PATCH] Improve `redis.cluster` types (#8573)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Alex Waygood <Alex.Waygood@Gmail.com>
---
 stubs/redis/redis/cluster.pyi | 66 +++++++++++++++++++----------------
 1 file changed, 35 insertions(+), 31 deletions(-)

diff --git a/stubs/redis/redis/cluster.pyi b/stubs/redis/redis/cluster.pyi
index c32379de6..325c24f23 100644
--- a/stubs/redis/redis/cluster.pyi
+++ b/stubs/redis/redis/cluster.pyi
@@ -1,5 +1,5 @@
 from _typeshed import Incomplete, Self
-from collections.abc import Callable
+from collections.abc import Callable, Iterable
 from threading import Lock
 from types import TracebackType
 from typing import Any, ClassVar, Generic
@@ -9,7 +9,7 @@ from redis.client import CaseInsensitiveDict, PubSub, Redis
 from redis.commands import CommandsParser, RedisClusterCommands
 from redis.commands.core import _StrType
 from redis.connection import BaseParser, Connection, Encoder
-from redis.exceptions import RedisError
+from redis.exceptions import MovedError, RedisError
 
 def get_node_name(host: str, port: str | int) -> str: ...
 def get_connection(redis_node: Redis[Any], *args, **options) -> Connection: ...
@@ -23,7 +23,7 @@ SLOT_ID: str
 REDIS_ALLOWED_KEYS: tuple[str, ...]
 KWARGS_DISABLED_KEYS: tuple[str, ...]
 
-def cleanup_kwargs(**kwargs): ...
+def cleanup_kwargs(**kwargs: Any) -> dict[str, Any]: ...
 
 # It uses `DefaultParser` in real life, but it is a dynamic base class.
 class ClusterParser(BaseParser): ...
@@ -106,47 +106,51 @@ class RedisCluster(AbstractRedisCluster, RedisClusterCommands[_StrType], Generic
     def close(self) -> None: ...
 
 class ClusterNode:
-    host: Any
-    port: Any
-    name: Any
-    server_type: Any
-    redis_connection: Any
-    def __init__(self, host, port, server_type: Any | None = ..., redis_connection: Any | None = ...) -> None: ...
-    def __eq__(self, obj): ...
+    host: str
+    port: int
+    name: str
+    server_type: str | None
+    redis_connection: Redis[Incomplete] | None
+    def __init__(
+        self, host: str, port: int, server_type: str | None = ..., redis_connection: Redis[Incomplete] | None = ...
+    ) -> None: ...
+    def __eq__(self, obj: object) -> bool: ...
     def __del__(self) -> None: ...
 
 class LoadBalancer:
-    primary_to_idx: Any
-    start_index: Any
+    primary_to_idx: dict[str, int]
+    start_index: int
     def __init__(self, start_index: int = ...) -> None: ...
-    def get_server_index(self, primary, list_size): ...
+    def get_server_index(self, primary: str, list_size: int) -> int: ...
     def reset(self) -> None: ...
 
 class NodesManager:
-    nodes_cache: Any
-    slots_cache: Any
-    startup_nodes: Any
-    default_node: Any
-    from_url: Any
-    connection_kwargs: Any
-    read_load_balancer: Any
+    nodes_cache: dict[str, ClusterNode]
+    slots_cache: dict[str, list[ClusterNode]]
+    startup_nodes: dict[str, ClusterNode]
+    default_node: ClusterNode | None
+    from_url: bool
+    connection_kwargs: dict[str, Any]
+    read_load_balancer: LoadBalancer
     def __init__(
         self,
-        startup_nodes,
+        startup_nodes: Iterable[ClusterNode],
         from_url: bool = ...,
         require_full_coverage: bool = ...,
-        lock: Incomplete | None = ...,
+        lock: Lock | None = ...,
         dynamic_startup_nodes: bool = ...,
-        **kwargs,
+        **kwargs: Any,
     ) -> None: ...
-    def get_node(self, host: Any | None = ..., port: Any | None = ..., node_name: Any | None = ...): ...
-    def update_moved_exception(self, exception) -> None: ...
-    def get_node_from_slot(self, slot, read_from_replicas: bool = ..., server_type: Any | None = ...): ...
-    def get_nodes_by_server_type(self, server_type): ...
-    def populate_startup_nodes(self, nodes) -> None: ...
-    def check_slots_coverage(self, slots_cache): ...
-    def create_redis_connections(self, nodes) -> None: ...
-    def create_redis_node(self, host, port, **kwargs): ...
+    def get_node(
+        self, host: str | None = ..., port: int | str | None = ..., node_name: str | None = ...
+    ) -> ClusterNode | None: ...
+    def update_moved_exception(self, exception: MovedError) -> None: ...
+    def get_node_from_slot(self, slot: str, read_from_replicas: bool = ..., server_type: str | None = ...) -> ClusterNode: ...
+    def get_nodes_by_server_type(self, server_type: str) -> list[ClusterNode]: ...
+    def populate_startup_nodes(self, nodes: Iterable[ClusterNode]) -> None: ...
+    def check_slots_coverage(self, slots_cache: dict[str, list[ClusterNode]]) -> bool: ...
+    def create_redis_connections(self, nodes: Iterable[ClusterNode]) -> None: ...
+    def create_redis_node(self, host: str, port: int | str, **kwargs: Any) -> Redis[Incomplete]: ...
     def initialize(self) -> None: ...
     def close(self) -> None: ...
     def reset(self) -> None: ...