TEST

Author: Takuma Iwaki

channels_redis/core.py

@@ -37,7 +37,7 @@ class ChannelLock:
     to mitigate multi-event loop problems.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         self.locks: collections.defaultdict[str, asyncio.Lock] = (
             collections.defaultdict(asyncio.Lock)
         )
@@ -58,7 +58,7 @@ def locked(self, channel: str) -> bool:
         """
         return self.locks[channel].locked()
 
-    def release(self, channel: str):
+    def release(self, channel: str) -> None:
         """
         Release the lock for the given channel.
         """
@@ -69,8 +69,8 @@ def release(self, channel: str):
             del self.wait_counts[channel]
 
 
-class BoundedQueue(asyncio.Queue):
-    def put_nowait(self, item):
+class BoundedQueue(asyncio.Queue[typing.Any]):
+    def put_nowait(self, item: typing.Any) -> None:
         if self.full():
             # see: https://github.com/django/channels_redis/issues/212
             # if we actually get into this code block, it likely means that
@@ -83,7 +83,7 @@ def put_nowait(self, item):
 
 
 class RedisLoopLayer:
-    def __init__(self, channel_layer: "RedisChannelLayer"):
+    def __init__(self, channel_layer: "RedisChannelLayer") -> None:
         self._lock = asyncio.Lock()
         self.channel_layer = channel_layer
         self._connections: typing.Dict[int, "Redis"] = {}
@@ -95,7 +95,7 @@ def get_connection(self, index: int) -> "Redis":
 
         return self._connections[index]
 
-    async def flush(self):
+    async def flush(self) -> None:
         async with self._lock:
             for index in list(self._connections):
                 connection = self._connections.pop(index)
@@ -116,15 +116,15 @@ class RedisChannelLayer(BaseChannelLayer):
     def __init__(
         self,
         hosts=None,
-        prefix="asgi",
+        prefix: str = "asgi",
         expiry=60,
-        group_expiry=86400,
+        group_expiry: int = 86400,
         capacity=100,
         channel_capacity=None,
         symmetric_encryption_keys=None,
         random_prefix_length=12,
         serializer_format="msgpack",
-    ):
+    ) -> None:
         # Store basic information
         self.expiry = expiry
         self.group_expiry = group_expiry
@@ -161,7 +161,7 @@ def __init__(
             collections.defaultdict(functools.partial(BoundedQueue, self.capacity))
         )
         # Detached channel cleanup tasks
-        self.receive_cleaners: typing.List[asyncio.Task] = []
+        self.receive_cleaners: typing.List[asyncio.Task[typing.Any]] = []
         # Per-channel cleanup locks to prevent a receive starting and moving
         # a message back into the main queue before its cleanup has completed
         self.receive_clean_locks = ChannelLock()
@@ -173,7 +173,7 @@ def create_pool(self, index: int) -> "ConnectionPool":
 
     extensions = ["groups", "flush"]
 
-    async def send(self, channel: str, message):
+    async def send(self, channel: str, message: typing.Any) -> None:
         """
         Send a message onto a (general or specific) channel.
         """
@@ -221,7 +221,7 @@ def _backup_channel_name(self, channel: str) -> str:
 
     async def _brpop_with_clean(
         self, index: int, channel: str, timeout: typing.Union[int, float, bytes, str]
-    ):
+    ) -> typing.Any:
         """
         Perform a Redis BRPOP and manage the backup processing queue.
         In case of cancellation, make sure the message is not lost.
@@ -252,15 +252,15 @@ async def _brpop_with_clean(
 
         return member
 
-    async def _clean_receive_backup(self, index: int, channel: str):
+    async def _clean_receive_backup(self, index: int, channel: str) -> None:
         """
         Pop the oldest message off the channel backup queue.
         The result isn't interesting as it was already processed.
         """
         connection = self.connection(index)
         await connection.zpopmin(self._backup_channel_name(channel))
 
-    async def receive(self, channel: str):
+    async def receive(self, channel: str) -> typing.Any:
         """
         Receive the first message that arrives on the channel.
         If more than one coroutine waits on the same channel, the first waiter
@@ -271,9 +271,9 @@ async def receive(self, channel: str):
         assert self.valid_channel_name(channel)
         if "!" in channel:
             real_channel = self.non_local_name(channel)
-            assert real_channel.endswith(
-                self.client_prefix + "!"
-            ), "Wrong client prefix"
+            assert real_channel.endswith(self.client_prefix + "!"), (
+                "Wrong client prefix"
+            )
             # Enter receiving section
             loop = asyncio.get_running_loop()
             self.receive_count += 1
@@ -292,11 +292,11 @@ async def receive(self, channel: str):
                 # Wait for our message to appear
                 message = None
                 while self.receive_buffer[channel].empty():
-                    tasks = [
+                    _tasks = [
                         self.receive_lock.acquire(),
                         self.receive_buffer[channel].get(),
                     ]
-                    tasks = [asyncio.ensure_future(task) for task in tasks]
+                    tasks = [asyncio.ensure_future(task) for task in _tasks]
                     try:
                         done, pending = await asyncio.wait(
                             tasks, return_when=asyncio.FIRST_COMPLETED
@@ -384,7 +384,9 @@ async def receive(self, channel: str):
             # Do a plain direct receive
             return (await self.receive_single(channel))[1]
 
-    async def receive_single(self, channel: str) -> typing.Tuple:
+    async def receive_single(
+        self, channel: str
+    ) -> typing.Tuple[typing.Any, typing.Any]:
         """
         Receives a single message off of the channel and returns it.
         """
@@ -420,7 +422,7 @@ async def receive_single(self, channel: str) -> typing.Tuple:
             )
             self.receive_cleaners.append(cleaner)
 
-            def _cleanup_done(cleaner: asyncio.Task):
+            def _cleanup_done(cleaner: asyncio.Task[typing.Any]) -> None:
                 self.receive_cleaners.remove(cleaner)
                 self.receive_clean_locks.release(channel_key)
 
@@ -448,7 +450,7 @@ async def new_channel(self, prefix: str = "specific") -> str:
 
     ### Flush extension ###
 
-    async def flush(self):
+    async def flush(self) -> None:
         """
         Deletes all messages and groups on all shards.
         """
@@ -470,7 +472,7 @@ async def flush(self):
         # Now clear the pools as well
         await self.close_pools()
 
-    async def close_pools(self):
+    async def close_pools(self) -> None:
         """
         Close all connections in the event loop pools.
         """
@@ -480,7 +482,7 @@ async def close_pools(self):
         for layer in self._layers.values():
             await layer.flush()
 
-    async def wait_received(self):
+    async def wait_received(self) -> None:
         """
         Wait for all channel cleanup functions to finish.
         """
@@ -489,13 +491,13 @@ async def wait_received(self):
 
     ### Groups extension ###
 
-    async def group_add(self, group: str, channel: str):
+    async def group_add(self, group: str, channel: str) -> None:
         """
         Adds the channel name to a group.
         """
         # Check the inputs
-        assert self.valid_group_name(group), True
-        assert self.valid_channel_name(channel), True
+        assert self.valid_group_name(group), "Group name not valid"
+        assert self.valid_channel_name(channel), "Channel name not valid"
         # Get a connection to the right shard
         group_key = self._group_key(group)
         connection = self.connection(self.consistent_hash(group))
@@ -505,7 +507,7 @@ async def group_add(self, group: str, channel: str):
         # it at this point is guaranteed to expire before that
         await connection.expire(group_key, self.group_expiry)
 
-    async def group_discard(self, group: str, channel: str):
+    async def group_discard(self, group: str, channel: str) -> None:
         """
         Removes the channel from the named group if it is in the group;
         does nothing otherwise (does not error)
@@ -516,7 +518,7 @@ async def group_discard(self, group: str, channel: str):
         connection = self.connection(self.consistent_hash(group))
         await connection.zrem(key, channel)
 
-    async def group_send(self, group: str, message):
+    async def group_send(self, group: str, message: typing.Any) -> None:
         """
         Sends a message to the entire group.
         """
@@ -540,9 +542,9 @@ async def group_send(self, group: str, message):
         for connection_index, channel_redis_keys in connection_to_channel_keys.items():
             # Discard old messages based on expiry
             pipe = connection.pipeline()
-            for key in channel_redis_keys:
+            for _key in channel_redis_keys:
                 pipe.zremrangebyscore(
-                    key, min=0, max=int(time.time()) - int(self.expiry)
+                    _key, min=0, max=int(time.time()) - int(self.expiry)
                 )
             await pipe.execute()
 
@@ -585,7 +587,7 @@ async def group_send(self, group: str, message):
             channels_over_capacity = await connection.eval(
                 group_send_lua, len(channel_redis_keys), *channel_redis_keys, *args
             )
-            _channels_over_capacity = -1
+            _channels_over_capacity = -1.0
             try:
                 _channels_over_capacity = float(channels_over_capacity)
             except Exception:
@@ -598,7 +600,13 @@ async def group_send(self, group: str, message):
                     group,
                 )
 
-    def _map_channel_keys_to_connection(self, channel_names, message):
+    def _map_channel_keys_to_connection(
+        self, channel_names: typing.Iterable[str], message: typing.Any
+    ) -> typing.Tuple[
+        typing.Dict[int, typing.List[str]],
+        typing.Dict[str, typing.Any],
+        typing.Dict[str, int],
+    ]:
         """
         For a list of channel names, GET
 
@@ -611,19 +619,21 @@ def _map_channel_keys_to_connection(self, channel_names, message):
         """
 
         # Connection dict keyed by index to list of redis keys mapped on that index
-        connection_to_channel_keys = collections.defaultdict(list)
+        connection_to_channel_keys: typing.Dict[int, typing.List[str]] = (
+            collections.defaultdict(list)
+        )
         # Message dict maps redis key to the message that needs to be send on that key
-        channel_key_to_message = dict()
+        channel_key_to_message: typing.Dict[str, typing.Any] = dict()
         # Channel key mapped to its capacity
-        channel_key_to_capacity = dict()
+        channel_key_to_capacity: typing.Dict[str, int] = dict()
 
         # For each channel
         for channel in channel_names:
             channel_non_local_name = channel
             if "!" in channel:
                 channel_non_local_name = self.non_local_name(channel)
             # Get its redis key
-            channel_key = self.prefix + channel_non_local_name
+            channel_key: str = self.prefix + channel_non_local_name
             # Have we come across the same redis key?
             if channel_key not in channel_key_to_message:
                 # If not, fill the corresponding dicts
@@ -654,13 +664,15 @@ def _group_key(self, group: str) -> bytes:
         """
         return f"{self.prefix}:group:{group}".encode("utf8")
 
-    def serialize(self, message) -> bytes:
+    ### Serialization ###
+
+    def serialize(self, message: typing.Any) -> bytes:
         """
         Serializes message to a byte string.
         """
         return self._serializer.serialize(message)
 
-    def deserialize(self, message: bytes):
+    def deserialize(self, message: bytes) -> typing.Any:
         """
         Deserializes from a byte string.
         """
@@ -671,7 +683,7 @@ def deserialize(self, message: bytes):
     def consistent_hash(self, value: typing.Union[str, "Buffer"]) -> int:
         return _consistent_hash(value, self.ring_size)
 
-    def __str__(self):
+    def __str__(self) -> str:
         return f"{self.__class__.__name__}(hosts={self.hosts})"
 
     ### Connection handling ###