API refactor

venv/lib/python3.12/site-packages/redis/_parsers/hiredis.py

@@ -1,19 +1,23 @@
 import asyncio
 import socket
 import sys
-from typing import Callable, List, Optional, Union
+from logging import getLogger
+from typing import Callable, List, Optional, TypedDict, Union
 
 if sys.version_info.major >= 3 and sys.version_info.minor >= 11:
     from asyncio import timeout as async_timeout
 else:
     from async_timeout import timeout as async_timeout
 
-from redis.compat import TypedDict
-
 from ..exceptions import ConnectionError, InvalidResponse, RedisError
 from ..typing import EncodableT
 from ..utils import HIREDIS_AVAILABLE
-from .base import AsyncBaseParser, BaseParser
+from .base import (
+    AsyncBaseParser,
+    AsyncPushNotificationsParser,
+    BaseParser,
+    PushNotificationsParser,
+)
 from .socket import (
     NONBLOCKING_EXCEPTION_ERROR_NUMBERS,
     NONBLOCKING_EXCEPTIONS,
@@ -21,6 +25,11 @@ from .socket import (
     SERVER_CLOSED_CONNECTION_ERROR,
 )
 
+# Used to signal that hiredis-py does not have enough data to parse.
+# Using `False` or `None` is not reliable, given that the parser can
+# return `False` or `None` for legitimate reasons from RESP payloads.
+NOT_ENOUGH_DATA = object()
+
 
 class _HiredisReaderArgs(TypedDict, total=False):
     protocolError: Callable[[str], Exception]
@@ -29,7 +38,7 @@ class _HiredisReaderArgs(TypedDict, total=False):
     errors: Optional[str]
 
 
-class _HiredisParser(BaseParser):
+class _HiredisParser(BaseParser, PushNotificationsParser):
     "Parser class for connections using Hiredis"
 
     def __init__(self, socket_read_size):
@@ -37,6 +46,9 @@ class _HiredisParser(BaseParser):
             raise RedisError("Hiredis is not installed")
         self.socket_read_size = socket_read_size
         self._buffer = bytearray(socket_read_size)
+        self.pubsub_push_handler_func = self.handle_pubsub_push_response
+        self.invalidation_push_handler_func = None
+        self._hiredis_PushNotificationType = None
 
     def __del__(self):
         try:
@@ -44,6 +56,11 @@ class _HiredisParser(BaseParser):
         except Exception:
             pass
 
+    def handle_pubsub_push_response(self, response):
+        logger = getLogger("push_response")
+        logger.debug("Push response: " + str(response))
+        return response
+
     def on_connect(self, connection, **kwargs):
         import hiredis
 
@@ -53,25 +70,32 @@ class _HiredisParser(BaseParser):
             "protocolError": InvalidResponse,
             "replyError": self.parse_error,
             "errors": connection.encoder.encoding_errors,
+            "notEnoughData": NOT_ENOUGH_DATA,
         }
 
         if connection.encoder.decode_responses:
             kwargs["encoding"] = connection.encoder.encoding
         self._reader = hiredis.Reader(**kwargs)
-        self._next_response = False
+        self._next_response = NOT_ENOUGH_DATA
+
+        try:
+            self._hiredis_PushNotificationType = hiredis.PushNotification
+        except AttributeError:
+            # hiredis < 3.2
+            self._hiredis_PushNotificationType = None
 
     def on_disconnect(self):
         self._sock = None
         self._reader = None
-        self._next_response = False
+        self._next_response = NOT_ENOUGH_DATA
 
     def can_read(self, timeout):
         if not self._reader:
             raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
 
-        if self._next_response is False:
+        if self._next_response is NOT_ENOUGH_DATA:
             self._next_response = self._reader.gets()
-            if self._next_response is False:
+            if self._next_response is NOT_ENOUGH_DATA:
                 return self.read_from_socket(timeout=timeout, raise_on_timeout=False)
         return True
 
@@ -105,14 +129,24 @@ class _HiredisParser(BaseParser):
             if custom_timeout:
                 sock.settimeout(self._socket_timeout)
 
-    def read_response(self, disable_decoding=False):
+    def read_response(self, disable_decoding=False, push_request=False):
         if not self._reader:
             raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
 
         # _next_response might be cached from a can_read() call
-        if self._next_response is not False:
+        if self._next_response is not NOT_ENOUGH_DATA:
             response = self._next_response
-            self._next_response = False
+            self._next_response = NOT_ENOUGH_DATA
+            if self._hiredis_PushNotificationType is not None and isinstance(
+                response, self._hiredis_PushNotificationType
+            ):
+                response = self.handle_push_response(response)
+                if not push_request:
+                    return self.read_response(
+                        disable_decoding=disable_decoding, push_request=push_request
+                    )
+                else:
+                    return response
             return response
 
         if disable_decoding:
@@ -120,7 +154,7 @@ class _HiredisParser(BaseParser):
         else:
             response = self._reader.gets()
 
-        while response is False:
+        while response is NOT_ENOUGH_DATA:
             self.read_from_socket()
             if disable_decoding:
                 response = self._reader.gets(False)
@@ -131,6 +165,16 @@ class _HiredisParser(BaseParser):
         # happened
         if isinstance(response, ConnectionError):
             raise response
+        elif self._hiredis_PushNotificationType is not None and isinstance(
+            response, self._hiredis_PushNotificationType
+        ):
+            response = self.handle_push_response(response)
+            if not push_request:
+                return self.read_response(
+                    disable_decoding=disable_decoding, push_request=push_request
+                )
+            else:
+                return response
         elif (
             isinstance(response, list)
             and response
@@ -140,7 +184,7 @@ class _HiredisParser(BaseParser):
         return response
 
 
-class _AsyncHiredisParser(AsyncBaseParser):
+class _AsyncHiredisParser(AsyncBaseParser, AsyncPushNotificationsParser):
     """Async implementation of parser class for connections using Hiredis"""
 
     __slots__ = ("_reader",)
@@ -150,6 +194,14 @@ class _AsyncHiredisParser(AsyncBaseParser):
             raise RedisError("Hiredis is not available.")
         super().__init__(socket_read_size=socket_read_size)
         self._reader = None
+        self.pubsub_push_handler_func = self.handle_pubsub_push_response
+        self.invalidation_push_handler_func = None
+        self._hiredis_PushNotificationType = None
+
+    async def handle_pubsub_push_response(self, response):
+        logger = getLogger("push_response")
+        logger.debug("Push response: " + str(response))
+        return response
 
     def on_connect(self, connection):
         import hiredis
@@ -158,6 +210,7 @@ class _AsyncHiredisParser(AsyncBaseParser):
         kwargs: _HiredisReaderArgs = {
             "protocolError": InvalidResponse,
             "replyError": self.parse_error,
+            "notEnoughData": NOT_ENOUGH_DATA,
         }
         if connection.encoder.decode_responses:
             kwargs["encoding"] = connection.encoder.encoding
@@ -166,13 +219,21 @@ class _AsyncHiredisParser(AsyncBaseParser):
         self._reader = hiredis.Reader(**kwargs)
         self._connected = True
 
+        try:
+            self._hiredis_PushNotificationType = getattr(
+                hiredis, "PushNotification", None
+            )
+        except AttributeError:
+            # hiredis < 3.2
+            self._hiredis_PushNotificationType = None
+
     def on_disconnect(self):
         self._connected = False
 
     async def can_read_destructive(self):
         if not self._connected:
             raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
-        if self._reader.gets():
+        if self._reader.gets() is not NOT_ENOUGH_DATA:
             return True
         try:
             async with async_timeout(0):
@@ -190,7 +251,7 @@ class _AsyncHiredisParser(AsyncBaseParser):
         return True
 
     async def read_response(
-        self, disable_decoding: bool = False
+        self, disable_decoding: bool = False, push_request: bool = False
     ) -> Union[EncodableT, List[EncodableT]]:
         # If `on_disconnect()` has been called, prohibit any more reads
         # even if they could happen because data might be present.
@@ -198,16 +259,33 @@ class _AsyncHiredisParser(AsyncBaseParser):
         if not self._connected:
             raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR) from None
 
-        response = self._reader.gets()
-        while response is False:
-            await self.read_from_socket()
+        if disable_decoding:
+            response = self._reader.gets(False)
+        else:
             response = self._reader.gets()
 
+        while response is NOT_ENOUGH_DATA:
+            await self.read_from_socket()
+            if disable_decoding:
+                response = self._reader.gets(False)
+            else:
+                response = self._reader.gets()
+
         # if the response is a ConnectionError or the response is a list and
         # the first item is a ConnectionError, raise it as something bad
         # happened
         if isinstance(response, ConnectionError):
             raise response
+        elif self._hiredis_PushNotificationType is not None and isinstance(
+            response, self._hiredis_PushNotificationType
+        ):
+            response = await self.handle_push_response(response)
+            if not push_request:
+                return await self.read_response(
+                    disable_decoding=disable_decoding, push_request=push_request
+                )
+            else:
+                return response
         elif (
             isinstance(response, list)
             and response