2022-06-24 17:14:37 +02:00

110 lines
3.5 KiB
Python

import json
from aioredis.util import wait_make_dict
class PubSubCommandsMixin:
"""Pub/Sub commands mixin.
For commands details see: http://redis.io/commands/#pubsub
"""
def publish(self, channel, message):
"""Post a message to channel."""
return self.execute(b'PUBLISH', channel, message)
def publish_json(self, channel, obj):
"""Post a JSON-encoded message to channel."""
return self.publish(channel, json.dumps(obj))
def subscribe(self, channel, *channels):
"""Switch connection to Pub/Sub mode and
subscribe to specified channels.
Arguments can be instances of :class:`~aioredis.Channel`.
Returns :func:`asyncio.gather()` coroutine which when done will return
a list of :class:`~aioredis.Channel` objects.
"""
conn = self._pool_or_conn
return wait_return_channels(
conn.execute_pubsub(b'SUBSCRIBE', channel, *channels),
conn, 'pubsub_channels')
def unsubscribe(self, channel, *channels):
"""Unsubscribe from specific channels.
Arguments can be instances of :class:`~aioredis.Channel`.
"""
conn = self._pool_or_conn
return conn.execute_pubsub(b'UNSUBSCRIBE', channel, *channels)
def psubscribe(self, pattern, *patterns):
"""Switch connection to Pub/Sub mode and
subscribe to specified patterns.
Arguments can be instances of :class:`~aioredis.Channel`.
Returns :func:`asyncio.gather()` coroutine which when done will return
a list of subscribed :class:`~aioredis.Channel` objects with
``is_pattern`` property set to ``True``.
"""
conn = self._pool_or_conn
return wait_return_channels(
conn.execute_pubsub(b'PSUBSCRIBE', pattern, *patterns),
conn, 'pubsub_patterns')
def punsubscribe(self, pattern, *patterns):
"""Unsubscribe from specific patterns.
Arguments can be instances of :class:`~aioredis.Channel`.
"""
conn = self._pool_or_conn
return conn.execute_pubsub(b'PUNSUBSCRIBE', pattern, *patterns)
def pubsub_channels(self, pattern=None):
"""Lists the currently active channels."""
args = [b'PUBSUB', b'CHANNELS']
if pattern is not None:
args.append(pattern)
return self.execute(*args)
def pubsub_numsub(self, *channels):
"""Returns the number of subscribers for the specified channels."""
return wait_make_dict(self.execute(
b'PUBSUB', b'NUMSUB', *channels))
def pubsub_numpat(self):
"""Returns the number of subscriptions to patterns."""
return self.execute(b'PUBSUB', b'NUMPAT')
@property
def channels(self):
"""Returns read-only channels dict.
See :attr:`~aioredis.RedisConnection.pubsub_channels`
"""
return self._pool_or_conn.pubsub_channels
@property
def patterns(self):
"""Returns read-only patterns dict.
See :attr:`~aioredis.RedisConnection.pubsub_patterns`
"""
return self._pool_or_conn.pubsub_patterns
@property
def in_pubsub(self):
"""Indicates that connection is in PUB/SUB mode.
Provides the number of subscribed channels.
"""
return self._pool_or_conn.in_pubsub
async def wait_return_channels(fut, conn, field):
res = await fut
channels_dict = getattr(conn, field)
return [channels_dict[name] for cmd, name, count in res]