More rename, add few types

Author: Rafiot

valkey/cluster.pyi

@@ -5,16 +5,16 @@ from types import TracebackType
 from typing import Any, ClassVar, Literal, NoReturn, Protocol
 from typing_extensions import Self
 
-from redis.client import CaseInsensitiveDict, PubSub, Valkey, _ParseResponseOptions
-from redis.commands import CommandsParser, ValkeyClusterCommands
-from redis.commands.core import _StrType
-from redis.connection import BaseParser, Connection, ConnectionPool, Encoder, _ConnectionPoolOptions, _Encodable
-from redis.exceptions import MovedError, ValkeyError
-from redis.retry import Retry
-from redis.typing import EncodableT
+from valkey.client import CaseInsensitiveDict, PubSub, Valkey, _ParseResponseOptions
+from valkey.commands import CommandsParser, ValkeyClusterCommands
+from valkey.commands.core import _StrType
+from valkey.connection import BaseParser, Connection, ConnectionPool, Encoder, _ConnectionPoolOptions, _Encodable
+from valkey.exceptions import MovedError, ValkeyError
+from valkey.retry import Retry
+from valkey.typing import EncodableT
 
 def get_node_name(host: str, port: str | int) -> str: ...
-def get_connection(redis_node: Valkey[Any], *args, **options: _ConnectionPoolOptions) -> Connection: ...
+def get_connection(valkey_node: Valkey[Any], *args, **options: _ConnectionPoolOptions) -> Connection: ...
 def parse_scan_result(command: Unused, res, **options): ...
 def parse_pubsub_numsub(command: Unused, res, **options: Unused): ...
 def parse_cluster_slots(resp, **options) -> dict[tuple[int, int], dict[str, Any]]: ...
@@ -82,7 +82,7 @@ class ValkeyCluster(AbstractValkeyCluster, ValkeyClusterCommands[_StrType]):
     @classmethod
     def from_url(cls, url: str, **kwargs) -> Self: ...
     def on_connect(self, connection: Connection) -> None: ...
-    def get_redis_connection(self, node: ClusterNode) -> Valkey[Any]: ...
+    def get_valkey_connection(self, node: ClusterNode) -> Valkey[Any]: ...
     def get_node(
         self, host: str | None = None, port: str | int | None = None, node_name: str | None = None
     ) -> ClusterNode | None: ...
@@ -120,9 +120,9 @@ class ClusterNode:
     port: int
     name: str
     server_type: str | None
-    redis_connection: Valkey[Incomplete] | None
+    valkey_connection: Valkey[Incomplete] | None
     def __init__(
-        self, host: str, port: int, server_type: str | None = None, redis_connection: Valkey[Incomplete] | None = None
+        self, host: str, port: int, server_type: str | None = None, valkey_connection: Valkey[Incomplete] | None = None
     ) -> None: ...
     def __eq__(self, obj: object) -> bool: ...
     def __del__(self) -> None: ...
@@ -163,8 +163,8 @@ class NodesManager:
     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) -> Valkey[Incomplete]: ...
+    def create_valkey_connections(self, nodes: Iterable[ClusterNode]) -> None: ...
+    def create_valkey_node(self, host: str, port: int | str, **kwargs: Any) -> Valkey[Incomplete]: ...
     def initialize(self) -> None: ...
     def close(self) -> None: ...
     def reset(self) -> None: ...
@@ -175,7 +175,7 @@ class ClusterPubSub(PubSub):
     cluster: ValkeyCluster[Any]
     def __init__(
         self,
-        redis_cluster: ValkeyCluster[Any],
+        valkey_cluster: ValkeyCluster[Any],
         node: ClusterNode | None = None,
         host: str | None = None,
         port: int | None = None,
@@ -186,7 +186,7 @@ class ClusterPubSub(PubSub):
     ) -> None: ...
     def get_pubsub_node(self) -> ClusterNode | None: ...
     def execute_command(self, *args, **kwargs) -> None: ...
-    def get_redis_connection(self) -> Valkey[Any] | None: ...
+    def get_valkey_connection(self) -> Valkey[Any] | None: ...
 
 class ClusterPipeline(ValkeyCluster[_StrType]):
     command_stack: list[Incomplete]

valkey/commands/graph/__init__.pyi

@@ -24,3 +24,22 @@ class Graph(GraphCommands):
     def labels(self): ...
     def relationship_types(self): ...
     def property_keys(self): ...
+
+
+class AsyncGraph(GraphCommands):
+    NAME: Any
+    client: Any
+    execute_command: Any
+    nodes: Any
+    edges: Any
+    version: int
+    def __init__(self, client, name=...) -> None: ...
+    async def get_label(self, idx): ...
+    async def get_relation(self, idx): ...
+    async def get_property(self, idx): ...
+    async def add_node(self, node) -> None: ...
+    async def add_edge(self, edge) -> None: ...
+    async def call_procedure(self, procedure, *args, read_only: bool = False, **kwagrs): ...
+    async def labels(self): ...
+    async def relationship_types(self): ...
+    async def property_keys(self): ...

valkey/commands/graph/execution_plan.py

@@ -166,10 +166,10 @@ def _create_operation(args):
             args.pop(0)
             if len(args) > 0 and "Records produced" in args[-1]:
                 records_produced = int(
-                    re.search("Records produced: (\\d+)", args[-1]).group(1)
+                    re.search("Records produced: (\\d+)", args[-1]).group(1)  # type: ignore[union-attr]
                 )
                 execution_time = float(
-                    re.search("Execution time: (\\d+.\\d+) ms", args[-1]).group(1)
+                    re.search("Execution time: (\\d+.\\d+) ms", args[-1]).group(1)  # type: ignore[union-attr]
                 )
                 profile_stats = ProfileStats(records_produced, execution_time)
                 args.pop(-1)
@@ -186,15 +186,15 @@ def _create_operation(args):
                 # set the current operation and move next
                 child = _create_operation(current_op.split("|"))
                 if current:
-                    current = stack.pop()
+                    current = stack.pop()  # type: ignore[unreachable]
                     current.append_child(child)
                 current = child
                 i += 1
             elif op_level == level + 1:
                 # if the operation is child of the current operation
                 # add it as child and set as current operation
                 child = _create_operation(current_op.split("|"))
-                current.append_child(child)
+                current.append_child(child)  # type: ignore[union-attr]
                 stack.append(current)
                 current = child
                 level += 1

valkey/commands/search/__init__.pyi

@@ -20,3 +20,21 @@ class Search(SearchCommands):
         def commit(self): ...
 
     def __init__(self, client, index_name: str = "idx") -> None: ...
+
+class AsyncSearch(SearchCommands):
+    class BatchIndexer:
+        def __init__(self, client, chunk_size: int = 1000) -> None: ...
+        async def add_document(
+            self,
+            doc_id,
+            nosave: bool = False,
+            score: float = 1.0,
+            payload: Incomplete | None = None,
+            replace: bool = False,
+            partial: bool = False,
+            no_create: bool = False,
+            **fields,
+        ): ...
+        async def commit(self): ...
+
+    def __init__(self, client, index_name: str = "idx") -> None: ...

valkey/commands/search/aggregation.py

@@ -22,7 +22,7 @@ class Reducer:
     See the `valkeyearch.reducers` module for the actual reducers.
     """
 
-    NAME = None
+    NAME: Union[str, None] = None
 
     def __init__(self, *args: List[str]) -> None:
         self._args = args

valkey/commands/search/aggregation.pyi

@@ -1,4 +1,4 @@
-from typing import Any, ClassVar, Literal
+from typing import Any, ClassVar, Literal, Union
 
 FIELDNAME: Any
 
@@ -9,7 +9,7 @@ class Limit:
     def build_args(self): ...
 
 class Reducer:
-    NAME: ClassVar[None]
+    NAME: ClassVar[Union[str, None]]
     def __init__(self, *args) -> None: ...
     def alias(self, alias): ...
     @property

valkey/commands/search/field.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Union
 
 from valkey import DataError
 
@@ -18,10 +18,10 @@ class Field:
     def __init__(
         self,
         name: str,
-        args: List[str] = None,
+        args: Union[List[str], None] = None,
         sortable: bool = False,
         no_index: bool = False,
-        as_name: str = None,
+        as_name: Union[str, None] = None,
     ):
         if args is None:
             args = []
@@ -63,11 +63,11 @@ def __init__(
         name: str,
         weight: float = 1.0,
         no_stem: bool = False,
-        phonetic_matcher: str = None,
+        phonetic_matcher: Union[str, None] = None,
         withsuffixtrie: bool = False,
         **kwargs,
     ):
-        Field.__init__(self, name, args=[Field.TEXT, Field.WEIGHT, weight], **kwargs)
+        Field.__init__(self, name, args=[Field.TEXT, Field.WEIGHT, weight], **kwargs)  # type: ignore[list-item]
 
         if no_stem:
             Field.append_arg(self, self.NOSTEM)
@@ -180,5 +180,5 @@ def __init__(self, name: str, algorithm: str, attributes: dict, **kwargs):
             attr_li.extend([key, value])
 
         Field.__init__(
-            self, name, args=[Field.VECTOR, algorithm, len(attr_li), *attr_li], **kwargs
+            self, name, args=[Field.VECTOR, algorithm, len(attr_li), *attr_li], **kwargs  # type: ignore[list-item]
         )

valkey/commands/search/querystring.py

@@ -1,3 +1,6 @@
+from typing import Dict, List
+
+
 def tags(*t):
     """
     Indicate that the values should be matched to a tag field
@@ -182,7 +185,7 @@ def __init__(self, *children, **kwparams):
 
         self.params = []
 
-        kvparams = {}
+        kvparams: Dict[str, List[Value]] = {}
         for k, v in kwparams.items():
             curvals = kvparams.setdefault(k, [])
             if isinstance(v, (str, int, float)):

valkey/exceptions.pyi

@@ -10,6 +10,7 @@ class DataError(ValkeyError): ...
 class PubSubError(ValkeyError): ...
 class WatchError(ValkeyError): ...
 class NoScriptError(ResponseError): ...
+class OutOfMemoryError(ResponseError): ...
 class ExecAbortError(ResponseError): ...
 class ReadOnlyError(ResponseError): ...
 class NoPermissionError(ResponseError): ...

valkey/retry.py

@@ -7,7 +7,7 @@
 T = TypeVar("T")
 
 if TYPE_CHECKING:
-    from redis.backoff import AbstractBackoff
+    from valkey.backoff import AbstractBackoff
 
 
 class Retry: