async_conn: optimize

Author: utkarshgupta137

redis/asyncio/connection.py

@@ -24,7 +24,6 @@
     Type,
     TypeVar,
     Union,
-    cast,
 )
 from urllib.parse import ParseResult, parse_qs, unquote, urlparse
 
@@ -110,32 +109,32 @@ def __init__(self, encoding: str, encoding_errors: str, decode_responses: bool):
 
     def encode(self, value: EncodableT) -> EncodedT:
         """Return a bytestring or bytes-like representation of the value"""
+        if isinstance(value, str):
+            return value.encode(self.encoding, self.encoding_errors)
         if isinstance(value, (bytes, memoryview)):
             return value
-        if isinstance(value, bool):
-            # special case bool since it is a subclass of int
-            raise DataError(
-                "Invalid input of type: 'bool'. "
-                "Convert to a bytes, string, int or float first."
-            )
         if isinstance(value, (int, float)):
+            if isinstance(value, bool):
+                # special case bool since it is a subclass of int
+                raise DataError(
+                    "Invalid input of type: 'bool'. "
+                    "Convert to a bytes, string, int or float first."
+                )
             return repr(value).encode()
-        if not isinstance(value, str):
-            # a value we don't know how to deal with. throw an error
-            typename = value.__class__.__name__  # type: ignore[unreachable]
-            raise DataError(
-                f"Invalid input of type: {typename!r}. "
-                "Convert to a bytes, string, int or float first."
-            )
-        return value.encode(self.encoding, self.encoding_errors)
+        # a value we don't know how to deal with. throw an error
+        typename = value.__class__.__name__
+        raise DataError(
+            f"Invalid input of type: {typename!r}. "
+            "Convert to a bytes, string, int or float first."
+        )
 
     def decode(self, value: EncodableT, force=False) -> EncodableT:
         """Return a unicode string from the bytes-like representation"""
         if self.decode_responses or force:
-            if isinstance(value, memoryview):
-                return value.tobytes().decode(self.encoding, self.encoding_errors)
             if isinstance(value, bytes):
                 return value.decode(self.encoding, self.encoding_errors)
+            if isinstance(value, memoryview):
+                return value.tobytes().decode(self.encoding, self.encoding_errors)
         return value
 
 
@@ -336,7 +335,7 @@ def purge(self):
     def close(self):
         try:
             self.purge()
-            self._buffer.close()  # type: ignore[union-attr]
+            self._buffer.close()
         except Exception:
             # issue #633 suggests the purge/close somehow raised a
             # BadFileDescriptor error. Perhaps the client ran out of
@@ -466,7 +465,7 @@ def on_disconnect(self):
         self._next_response = False
 
     async def can_read(self, timeout: float):
-        if not self._reader:
+        if not self._stream or not self._reader:
             raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR)
 
         if self._next_response is False:
@@ -480,14 +479,14 @@ async def read_from_socket(
         timeout: Union[float, None, _Sentinel] = SENTINEL,
         raise_on_timeout: bool = True,
     ):
-        if self._stream is None or self._reader is None:
-            raise RedisError("Parser already closed.")
-
         timeout = self._socket_timeout if timeout is SENTINEL else timeout
         try:
-            async with async_timeout.timeout(timeout):
+            if timeout is None:
                 buffer = await self._stream.read(self._read_size)
-            if not isinstance(buffer, bytes) or len(buffer) == 0:
+            else:
+                async with async_timeout.timeout(timeout):
+                    buffer = await self._stream.read(self._read_size)
+            if not isinstance(buffer, bytes) or not buffer:
                 raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR) from None
             self._reader.feed(buffer)
             # data was read from the socket and added to the buffer.
@@ -516,9 +515,6 @@ async def read_response(
             self.on_disconnect()
             raise ConnectionError(SERVER_CLOSED_CONNECTION_ERROR) from None
 
-        response: Union[
-            EncodableT, ConnectionError, List[Union[EncodableT, ConnectionError]]
-        ]
         # _next_response might be cached from a can_read() call
         if self._next_response is not False:
             response = self._next_response
@@ -541,8 +537,7 @@ async def read_response(
             and isinstance(response[0], ConnectionError)
         ):
             raise response[0]
-        # cast as there won't be a ConnectionError here.
-        return cast(Union[EncodableT, List[EncodableT]], response)
+        return response
 
 
 DefaultParser: Type[Union[PythonParser, HiredisParser]]
@@ -637,7 +632,7 @@ def __init__(
         self.socket_type = socket_type
         self.retry_on_timeout = retry_on_timeout
         if retry_on_timeout:
-            if retry is None:
+            if not retry:
                 self.retry = Retry(NoBackoff(), 1)
             else:
                 # deep-copy the Retry object as it is mutable
@@ -681,7 +676,7 @@ def __del__(self):
 
     @property
     def is_connected(self):
-        return bool(self._reader and self._writer)
+        return self._reader and self._writer
 
     def register_connect_callback(self, callback):
         self._connect_callbacks.append(weakref.WeakMethod(callback))
@@ -713,7 +708,7 @@ async def connect(self):
             raise ConnectionError(exc) from exc
 
         try:
-            if self.redis_connect_func is None:
+            if not self.redis_connect_func:
                 # Use the default on_connect function
                 await self.on_connect()
             else:
@@ -745,7 +740,7 @@ async def _connect(self):
         self._reader = reader
         self._writer = writer
         sock = writer.transport.get_extra_info("socket")
-        if sock is not None:
+        if sock:
             sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
             try:
                 # TCP_KEEPALIVE
@@ -856,32 +851,29 @@ async def check_health(self):
             await self.retry.call_with_retry(self._send_ping, self._ping_failed)
 
     async def _send_packed_command(self, command: Iterable[bytes]) -> None:
-        if self._writer is None:
-            raise RedisError("Connection already closed.")
-
         self._writer.writelines(command)
         await self._writer.drain()
 
     async def send_packed_command(
-        self,
-        command: Union[bytes, str, Iterable[bytes]],
-        check_health: bool = True,
-    ):
-        """Send an already packed command to the Redis server"""
-        if not self._writer:
+        self, command: Union[bytes, str, Iterable[bytes]], check_health: bool = True
+    ) -> None:
+        if not self.is_connected:
             await self.connect()
-        # guard against health check recursion
-        if check_health:
+        elif check_health:
             await self.check_health()
+
         try:
             if isinstance(command, str):
                 command = command.encode()
             if isinstance(command, bytes):
                 command = [command]
-            await asyncio.wait_for(
-                self._send_packed_command(command),
-                self.socket_timeout,
-            )
+            if self.socket_timeout:
+                await asyncio.wait_for(
+                    self._send_packed_command(command), self.socket_timeout
+                )
+            else:
+                self._writer.writelines(command)
+                await self._writer.drain()
         except asyncio.TimeoutError:
             await self.disconnect()
             raise TimeoutError("Timeout writing to socket") from None
@@ -901,8 +893,6 @@ async def send_packed_command(
 
     async def send_command(self, *args, **kwargs):
         """Pack and send a command to the Redis server"""
-        if not self.is_connected:
-            await self.connect()
         await self.send_packed_command(
             self.pack_command(*args), check_health=kwargs.get("check_health", True)
         )
@@ -917,7 +907,12 @@ async def read_response(self, disable_decoding: bool = False):
         """Read the response from a previously sent command"""
         try:
             async with self._lock:
-                async with async_timeout.timeout(self.socket_timeout):
+                if self.socket_timeout:
+                    async with async_timeout.timeout(self.socket_timeout):
+                        response = await self._parser.read_response(
+                            disable_decoding=disable_decoding
+                        )
+                else:
                     response = await self._parser.read_response(
                         disable_decoding=disable_decoding
                     )
@@ -1176,10 +1171,7 @@ def __init__(
         self._lock = asyncio.Lock()
 
     def repr_pieces(self) -> Iterable[Tuple[str, Union[str, int]]]:
-        pieces = [
-            ("path", self.path),
-            ("db", self.db),
-        ]
+        pieces = [("path", self.path), ("db", self.db)]
         if self.client_name:
             pieces.append(("client_name", self.client_name))
         return pieces
@@ -1248,12 +1240,11 @@ def parse_url(url: str) -> ConnectKwargs:
             parser = URL_QUERY_ARGUMENT_PARSERS.get(name)
             if parser:
                 try:
-                    # We can't type this.
-                    kwargs[name] = parser(value)  # type: ignore[misc]
+                    kwargs[name] = parser(value)
                 except (TypeError, ValueError):
                     raise ValueError(f"Invalid value for `{name}` in connection URL.")
             else:
-                kwargs[name] = value  # type: ignore[misc]
+                kwargs[name] = value
 
     if parsed.username:
         kwargs["username"] = unquote(parsed.username)
@@ -1358,7 +1349,7 @@ def __init__(
         max_connections: Optional[int] = None,
         **connection_kwargs,
     ):
-        max_connections = max_connections or 2 ** 31
+        max_connections = max_connections or 2**31
         if not isinstance(max_connections, int) or max_connections < 0:
             raise ValueError('"max_connections" must be a positive integer')
 

tests/test_asyncio/test_connection_pool.py

@@ -664,7 +664,7 @@ async def r(self, create_redis):
 
     def assert_interval_advanced(self, connection):
         diff = connection.next_health_check - asyncio.get_event_loop().time()
-        assert self.interval > diff > (self.interval - 1)
+        assert self.interval >= diff > (self.interval - 1)
 
     async def test_health_check_runs(self, r):
         if r.connection:

tests/test_asyncio/test_lock.py

@@ -114,7 +114,7 @@ async def test_blocking_timeout(self, r, event_loop):
         start = event_loop.time()
         assert not await lock2.acquire()
         # The elapsed duration should be less than the total blocking_timeout
-        assert bt > (event_loop.time() - start) > bt - sleep
+        assert bt >= (event_loop.time() - start) > bt - sleep
         await lock1.release()
 
     async def test_context_manager(self, r):