POC

Author: petyaslavova

redis/asyncio/client.py

@@ -13,6 +13,7 @@
     Dict,
     Iterable,
     List,
+    Literal,
     Mapping,
     MutableMapping,
     Optional,
@@ -129,6 +130,9 @@ class Redis(
     Connection object to talk to redis.
     """
 
+    # Type discrimination marker for @overload self-type pattern
+    _is_async_client: Literal[True] = True
+
     response_callbacks: MutableMapping[Union[str, bytes], ResponseCallbackT]
 
     @classmethod

redis/asyncio/cluster.py

@@ -16,6 +16,7 @@
     Dict,
     Generator,
     List,
+    Literal,
     Mapping,
     Optional,
     Set,
@@ -238,6 +239,9 @@ def from_url(cls, url: str, **kwargs: Any) -> "RedisCluster":
             kwargs["ssl"] = True
         return cls(**kwargs)
 
+    # Type discrimination marker for @overload self-type pattern
+    _is_async_client: Literal[True] = True
+
     __slots__ = (
         "_initialize",
         "_lock",

redis/client.py

@@ -9,6 +9,7 @@
     Callable,
     Dict,
     List,
+    Literal,
     Mapping,
     Optional,
     Set,
@@ -135,6 +136,9 @@ class Redis(RedisModuleCommands, CoreCommands, SentinelCommands):
     It is not safe to pass PubSub or Pipeline objects between threads.
     """
 
+    # Type discrimination marker for @overload self-type pattern
+    _is_async_client: Literal[False] = False
+
     @classmethod
     def from_url(cls, url: str, **kwargs) -> "Redis":
         """

redis/cluster.py

@@ -554,6 +554,9 @@ def replace_default_node(self, target_node: "ClusterNode" = None) -> None:
 class RedisCluster(
     AbstractRedisCluster, MaintNotificationsAbstractRedisCluster, RedisClusterCommands
 ):
+    # Type discrimination marker for @overload self-type pattern
+    _is_async_client: Literal[False] = False
+
     @classmethod
     def from_url(cls, url: str, **kwargs: Any) -> "RedisCluster":
         """

redis/commands/core.py

@@ -31,6 +31,7 @@
     Set,
     Tuple,
     Union,
+    overload,
 )
 
 from redis.asyncio.observability.recorder import (
@@ -40,6 +41,7 @@
 from redis.typing import (
     AbsExpiryT,
     AnyKeyT,
+    AsyncClientProtocol,
     BitfieldOffsetT,
     ChannelT,
     CommandsProtocol,
@@ -55,6 +57,7 @@
     ResponseT,
     ScriptTextT,
     StreamIdT,
+    SyncClientProtocol,
     TimeoutSecT,
     ZScoreBoundT,
 )
@@ -2064,6 +2067,15 @@ def digest(self, name: KeyT) -> Union[str, bytes, None]:
         # Bulk string response is already handled (bytes/str based on decode_responses)
         return self.execute_command("DIGEST", name)
 
+    # --- @overload pattern for get() ---
+    # Sync client returns bytes | None directly
+    @overload
+    def get(self: SyncClientProtocol, name: KeyT) -> Optional[bytes]: ...
+
+    # Async client returns Awaitable[bytes | None]
+    @overload
+    def get(self: AsyncClientProtocol, name: KeyT) -> Awaitable[Optional[bytes]]: ...
+
     def get(self, name: KeyT) -> ResponseT:
         """
         Return the value at key ``name``, or None if the key doesn't exist
@@ -2554,6 +2566,47 @@ def restore(
 
         return self.execute_command("RESTORE", *params)
 
+    # --- @overload pattern for set() ---
+    # Sync client returns bool | None directly
+    @overload
+    def set(
+        self: SyncClientProtocol,
+        name: KeyT,
+        value: EncodableT,
+        ex: Optional[ExpiryT] = ...,
+        px: Optional[ExpiryT] = ...,
+        nx: bool = ...,
+        xx: bool = ...,
+        keepttl: bool = ...,
+        get: bool = ...,
+        exat: Optional[AbsExpiryT] = ...,
+        pxat: Optional[AbsExpiryT] = ...,
+        ifeq: Optional[Union[bytes, str]] = ...,
+        ifne: Optional[Union[bytes, str]] = ...,
+        ifdeq: Optional[str] = ...,
+        ifdne: Optional[str] = ...,
+    ) -> Optional[bool]: ...
+
+    # Async client returns Awaitable[bool | None]
+    @overload
+    def set(
+        self: AsyncClientProtocol,
+        name: KeyT,
+        value: EncodableT,
+        ex: Optional[ExpiryT] = ...,
+        px: Optional[ExpiryT] = ...,
+        nx: bool = ...,
+        xx: bool = ...,
+        keepttl: bool = ...,
+        get: bool = ...,
+        exat: Optional[AbsExpiryT] = ...,
+        pxat: Optional[AbsExpiryT] = ...,
+        ifeq: Optional[Union[bytes, str]] = ...,
+        ifne: Optional[Union[bytes, str]] = ...,
+        ifdeq: Optional[str] = ...,
+        ifdne: Optional[str] = ...,
+    ) -> Awaitable[Optional[bool]]: ...
+
     @experimental_args(["ifeq", "ifne", "ifdeq", "ifdne"])
     def set(
         self,

redis/commands/redismodules.py

@@ -8,7 +8,7 @@
     from .json import JSON
     from .search import AsyncSearch, Search
     from .timeseries import TimeSeries
-    from .vectorset import VectorSet
+    from .vectorset import AsyncVectorSet, VectorSet
 
 
 class RedisModuleCommands:
@@ -99,3 +99,11 @@ def ft(self, index_name="idx") -> AsyncSearch:
 
         s = AsyncSearch(client=self, index_name=index_name)
         return s
+
+    def vset(self) -> AsyncVectorSet:
+        """Access the VectorSet commands namespace."""
+
+        from .vectorset import AsyncVectorSet
+
+        vset = AsyncVectorSet(client=self)
+        return vset

redis/commands/vectorset/__init__.py

@@ -1,4 +1,5 @@
 import json
+from typing import Literal
 
 from redis._parsers.helpers import pairs_to_dict
 from redis.commands.vectorset.utils import (
@@ -18,9 +19,11 @@
 )
 
 
-class VectorSet(VectorSetCommands):
+class _VectorSetBase(VectorSetCommands):
+    """Base class with shared initialization logic for VectorSet clients."""
+
     def __init__(self, client, **kwargs):
-        """Create a new VectorSet client."""
+        """Initialize VectorSet client with callbacks."""
         # Set the module commands' callbacks
         self._MODULE_CALLBACKS = {
             VEMB_CMD: parse_vemb_result,
@@ -44,3 +47,15 @@ def __init__(self, client, **kwargs):
 
         for k, v in self._MODULE_CALLBACKS.items():
             self.client.set_response_callback(k, v)
+
+
+class VectorSet(_VectorSetBase):
+    """Sync VectorSet client."""
+
+    _is_async_client: Literal[False] = False
+
+
+class AsyncVectorSet(_VectorSetBase):
+    """Async VectorSet client."""
+
+    _is_async_client: Literal[True] = True

redis/commands/vectorset/commands.py

@@ -1,11 +1,19 @@
 import json
 from enum import Enum
-from typing import Any, Awaitable, Dict, List, Optional, Union
+from typing import Any, Awaitable, Dict, List, Optional, Union, overload
 
 from redis.client import NEVER_DECODE
 from redis.commands.helpers import get_protocol_version
 from redis.exceptions import DataError
-from redis.typing import CommandsProtocol, EncodableT, KeyT, Number
+from redis.typing import (
+    AsyncClientProtocol,
+    CommandsProtocol,
+    EncodableT,
+    KeyT,
+    Number,
+    ResponseT,
+    SyncClientProtocol,
+)
 
 VADD_CMD = "VADD"
 VSIM_CMD = "VSIM"
@@ -129,6 +137,38 @@ def vadd(
 
         return self.execute_command(VADD_CMD, key, *pieces)
 
+    @overload
+    def vsim(
+        self: SyncClientProtocol,
+        key: KeyT,
+        input: Union[List[float], bytes, str],
+        with_scores: Optional[bool] = ...,
+        with_attribs: Optional[bool] = ...,
+        count: Optional[int] = ...,
+        ef: Optional[Number] = ...,
+        filter: Optional[str] = ...,
+        filter_ef: Optional[str] = ...,
+        truth: Optional[bool] = ...,
+        no_thread: Optional[bool] = ...,
+        epsilon: Optional[Number] = ...,
+    ) -> VSimResult: ...
+
+    @overload
+    def vsim(
+        self: AsyncClientProtocol,
+        key: KeyT,
+        input: Union[List[float], bytes, str],
+        with_scores: Optional[bool] = ...,
+        with_attribs: Optional[bool] = ...,
+        count: Optional[int] = ...,
+        ef: Optional[Number] = ...,
+        filter: Optional[str] = ...,
+        filter_ef: Optional[str] = ...,
+        truth: Optional[bool] = ...,
+        no_thread: Optional[bool] = ...,
+        epsilon: Optional[Number] = ...,
+    ) -> Awaitable[VSimResult]: ...
+
     def vsim(
         self,
         key: KeyT,

redis/typing.py

@@ -6,11 +6,13 @@
     Any,
     Awaitable,
     Iterable,
+    Literal,
     Mapping,
     Protocol,
     Type,
     TypeVar,
     Union,
+    runtime_checkable,
 )
 
 if TYPE_CHECKING:
@@ -19,6 +21,30 @@
 
 
 Number = Union[int, float]
+
+
+@runtime_checkable
+class AsyncClientProtocol(Protocol):
+    """Protocol for asynchronous Redis clients (redis.asyncio.client.Redis).
+
+    This protocol uses a Literal marker to identify async clients.
+    Used in @overload to provide correct return types for async clients.
+    """
+
+    _is_async_client: Literal[True]
+
+
+@runtime_checkable
+class SyncClientProtocol(Protocol):
+    """Protocol for synchronous Redis clients (redis.client.Redis).
+
+    This protocol uses a Literal marker to identify sync clients.
+    Used in @overload to provide correct return types for sync clients.
+    """
+
+    _is_async_client: Literal[False]
+
+
 EncodedT = Union[bytes, bytearray, memoryview]
 DecodedT = Union[str, int, float]
 EncodableT = Union[EncodedT, DecodedT]