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,24 +109,24 @@ 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"""
@@ -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,12 +465,12 @@ 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:
+ if not self._next_response:
  self._next_response = self._reader.gets()
- if self._next_response is False:
+ if not self._next_response:
  return await self.read_from_socket(timeout=timeout, raise_on_timeout=False)
  return True
 
@@ -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,11 +515,8 @@ 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:
+ if self._next_response:
  response = self._next_response
  self._next_response = False
  return 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)

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):