<commit_before>#!/usr/bin/env python3
import asyncio
import logging
import types

from asyncio.futures import Future
from asyncio.queues import Queue
from asyncio.streams import StreamReader

try:
    import hiredis
except ImportError:
    hiredis = None

from collections import deque
from functools import wraps
from inspect import getfullargspec, formatargspec, getcallargs

from .encoders import BaseEncoder, UTF8Encoder
from .exceptions import (
        ConnectionLostError,
        Error,
        ErrorReply,
        NoRunningScriptError,
        NotConnectedError,
        ScriptKilledError,
        TimeoutError,
        TransactionError,
)
from .log import logger
from .replies import (
        BlockingPopReply,
        ClientListReply,
        ConfigPairReply,
        DictReply,
        EvalScriptReply,
        InfoReply,
        ListReply,
        PubSubReply,
        SetReply,
        StatusReply,
        ZRangeReply,
)


from .cursors import Cursor, SetCursor, DictCursor, ZCursor

__all__ = (
    'RedisProtocol',
    'HiRedisProtocol',
    'Transaction',
    'Subscription',
    'Script',

    'ZAggregate',
    'ZScoreBoundary',
)

NoneType = type(None)


class ZScoreBoundary:
    """
    Score boundary for a sorted set.
    for queries like zrangebyscore and similar

    :param value: Value for the boundary.
    :type value: float
    :param exclude_boundary: Exclude the boundary.
    :type exclude_boundary: bool
    """
    def __init__(self, value, exclude_boundary=False):
        assert isinstance(value, float) or value in ('+inf', '-inf')
        self.value = value
        self.exclude_boundary = exclude_boundary

    def __repr__(self):
        return 'ZScoreBoundary(value=%r, exclude_boundary=%r)' % (
                    self.value, self.exclude_boundary)

ZScoreBoundary.MIN_VALUE = ZScoreBoundary('-inf')
ZScoreBoundary.MAX_VALUE = ZScoreBoundary('+inf')


class ZAggregate: # TODO: use the Python 3.4 enum type.
    """
    Aggregation method for zinterstore and zunionstore.
    """
    #: Sum aggregation.
    SUM = 'SUM'

    #: Min aggregation.
    MIN = 'MIN'

    #: Max aggregation.
    MAX = 'MAX'


class PipelinedCall:
    """ Track record for call that is being executed in a protocol. """
    __slots__ = ('cmd', 'is_blocking')

    def __init__(self, cmd, is_blocking):
        self.cmd = cmd
        self.is_blocking = is_blocking


class MultiBulkReply:
    """
    Container for a multi bulk reply.
    """
    def __init__(self, protocol, count, loop=None):
        self._loop = loop or asyncio.get_event_loop()

        #: Buffer of incoming, undelivered data, received from the parser.
        self._data_queue = []

        #: Incoming read queries.
        #: Contains (read_count, Future, decode_flag, one_only_flag) tuples.
        self._f_queue = deque()

        self.protocol = protocol
        self.count = int(count)

    def _feed_received(self, item):
        """
        Feed entry for the parser.
        """
        # Push received items on the queue
        self._data_queue.append(item)
        self._flush()

    def _flush(self):
        """
        Answer read queries when we have enough data in our multibulk reply.
        """
        # As long as we have more data in our queue then we require for a read
        # query -> answer queries.
        while self._f_queue and self._f_queue[0][0] <= len(self._data_queue):
            # Pop query.
            count, f, decode, one_only = self._f_queue.popleft()

            # Slice data buffer.
            data, self._data_queue = self._data_queue[:count], self._data_queue[count:]

            # When the decode flag is given, decode bytes to native types.
            if decode:
                data = [ self._decode(d) for d in data ]

            # When one_only flag has been given, don't return an array.
            if one_only:
                assert len(data) == 1
                f.set_result(data[0])
            else:
                f.set_result(data)

    def _decode(self, result):
        """ Decode bytes to native Python types. """
        if isinstance(result, (StatusReply, int, float, MultiBulkReply)):
            # Note that MultiBulkReplies can be nested. e.g. in the 'scan' operation.
            return result
        elif isinstance(result, bytes):
            return self.protocol.decode_to_native(result)
        elif result is None:
            return result
        else:
            raise AssertionError('Invalid type: %r' % type(result))

    def _read(self, decode=True, count=1, _one=False):
        """ Do read operation on the queue. Return future. """
        f = Future(loop=self.protocol._loop)
        self._f_queue.append((count, f, decode, _one))

        # If there is enough data on the queue, answer future immediately.
        self._flush()

        return f

    def iter_raw(self):
        """
        Iterate over all multi bulk packets. This yields futures that won't
        decode bytes yet.
        """
        for i in range(self.count):
            yield self._read(decode=False, _one=True)

    def __iter__(self):
        """
        Iterate over the reply. This yields coroutines of the decoded packets.
        It decodes bytes automatically using protocol.decode_to_native.
        """
        for i in range(self.count):
            yield self._read(_one=True)

    def __repr__(self):
        return 'MultiBulkReply(protocol=%r, count=%r)' % (self.protocol, self.count)


class _ScanPart:
    """ Internal: result chunk of a scan operation. """
    def __init__(self, new_cursor_pos, items):
        self.new_cursor_pos = new_cursor_pos
        self.items = items


class PostProcessors:
    """
    At the protocol level, we only know about a few basic classes; they
    include: bool, int, StatusReply, MultiBulkReply and bytes.
    This will return a postprocessor function that turns these into more
    meaningful objects.

    For some methods, we have several post processors. E.g. a list can be
    returned either as a ListReply (which has some special streaming
    functionality), but also as a Python list.
    """
    @classmethod
    def get_all(cls, return_type):
        """
        Return list of (suffix, return_type, post_processor)
        """
        default = cls.get_default(return_type)
        alternate = cls.get_alternate_post_processor(return_type)

        result = [ ('', return_type, default) ]
        if alternate:
            result.append(alternate)
        return result

    @classmethod
    def get_default(cls, return_type):
        """ Give post processor function for return type. """
        return {
                ListReply: cls.multibulk_as_list,
                SetReply: cls.multibulk_as_set,
                DictReply: cls.multibulk_as_dict,

                float: cls.bytes_to_float,
                (float, NoneType): cls.bytes_to_float_or_none,
                NativeType: cls.bytes_to_native,
                (NativeType, NoneType): cls.bytes_to_native_or_none,
                InfoReply: cls.bytes_to_info,
                ClientListReply: cls.bytes_to_clientlist,
                str: cls.bytes_to_str,
                bool: cls.int_to_bool,
                BlockingPopReply: cls.multibulk_as_blocking_pop_reply,
                ZRangeReply: cls.multibulk_as_zrangereply,

                StatusReply: cls.bytes_to_status_reply,
                (StatusReply, NoneType): cls.bytes_to_status_reply_or_none,
                int: None,
                (int, NoneType): None,
                ConfigPairReply: cls.multibulk_as_configpair,
                ListOf(bool): cls.multibulk_as_boolean_list,
                _ScanPart: cls.multibulk_as_scanpart,
                EvalScriptReply: cls.any_to_evalscript,

                NoneType: None,
        }[return_type]

    @classmethod
    def get_alternate_post_processor(cls, return_type):
        """ For list/set/dict. Create additional post processors that return
        python classes rather than ListReply/SetReply/DictReply """
        original_post_processor = cls.get_default(return_type)

        if return_type == ListReply:
            @asyncio.coroutine
            def as_list(protocol, result):
                result = yield from original_post_processor(protocol, result)
                return (yield from result.aslist())
            return '_aslist', list, as_list

        elif return_type == SetReply:
            @asyncio.coroutine
            def as_set(protocol, result):
                result = yield from original_post_processor(protocol, result)
                return (yield from result.asset())
            return '_asset', set, as_set

        elif return_type in (DictReply, ZRangeReply):
            @asyncio.coroutine
            def as_dict(protocol, result):
                result = yield from original_post_processor(protocol, result)
                return (yield from result.asdict())
            return '_asdict', dict, as_dict

    # === Post processor handlers below. ===

    @asyncio.coroutine
    def multibulk_as_list(protocol, result):
        assert isinstance(result, MultiBulkReply)
        return ListReply(result)

    @asyncio.coroutine
    def multibulk_as_boolean_list(protocol, result):
        # Turn the array of integers into booleans.
        assert isinstance(result, MultiBulkReply)
        values = yield from ListReply(result).aslist()
        return [ bool(v) for v in values ]

    @asyncio.coroutine
    def multibulk_as_set(protocol, result):
        assert isinstance(result, MultiBulkReply)
        return SetReply(result)

    @asyncio.coroutine
    def multibulk_as_dict(protocol, result):
        assert isinstance(result, MultiBulkReply)
        return DictReply(result)

    @asyncio.coroutine
    def multibulk_as_zrangereply(protocol, result):
        assert isinstance(result, MultiBulkReply)
        return ZRangeReply(result)

    @asyncio.coroutine
    def multibulk_as_blocking_pop_reply(protocol, result):
        if result is None:
            raise TimeoutError('Timeout in blocking pop')
        else:
            assert isinstance(result, MultiBulkReply)
            list_name, value = yield from ListReply(result).aslist()
            return BlockingPopReply(list_name, value)

    @asyncio.coroutine
    def multibulk_as_configpair(protocol, result):
        assert isinstance(result, MultiBulkReply)
        parameter, value = yield from ListReply(result).aslist()
        return ConfigPairReply(parameter, value)

    @asyncio.coroutine
    def multibulk_as_scanpart(protocol, result):
        """
        Process scanpart result.
        This is a multibulk reply of length two, where the first item is the
        new cursor position and the second item is a nested multi bulk reply
        containing all the elements.
        """
        # Get outer multi bulk reply.
        assert isinstance(result, MultiBulkReply)
        new_cursor_pos, items_bulk = yield from ListReply(result).aslist()
        assert isinstance(items_bulk, MultiBulkReply)

        # Read all items for scan chunk in memory. This is fine, because it's
        # transmitted in chunks of about 10.
        items = yield from ListReply(items_bulk).aslist()
        return _ScanPart(int(new_cursor_pos), items)

    @asyncio.coroutine
    def bytes_to_info(protocol, result):
        assert isinstance(result, bytes)
        return InfoReply(result)

    @asyncio.coroutine
    def bytes_to_status_reply(protocol, result):
        assert isinstance(result, bytes)
        return StatusReply(result.decode('utf-8'))

    @asyncio.coroutine
    def bytes_to_status_reply_or_none(protocol, result):
        assert isinstance(result, (bytes, NoneType))
        if result:
            return StatusReply(result.decode('utf-8'))

    @asyncio.coroutine
    def bytes_to_clientlist(protocol, result):
        assert isinstance(result, bytes)
        return ClientListReply(result)

    @asyncio.coroutine
    def int_to_bool(protocol, result):
        assert isinstance(result, int)
        return bool(result) # Convert int to bool

    @asyncio.coroutine
    def bytes_to_native(protocol, result):
        assert isinstance(result, bytes)
        return protocol.decode_to_native(result)

    @asyncio.coroutine
    def bytes_to_str(protocol, result):
        assert isinstance(result, bytes)
        return result.decode('ascii')

    @asyncio.coroutine
    def bytes_to_native_or_none(protocol, result):
        if result is None:
            return result
        else:
            assert isinstance(result, bytes)
            return protocol.decode_to_native(result)

    @asyncio.coroutine
    def bytes_to_float_or_none(protocol, result):
        if result is None:
            return result
        assert isinstance(result, bytes)
        return float(result)

    @asyncio.coroutine
    def bytes_to_float(protocol, result):
        assert isinstance(result, bytes)
        return float(result)

    @asyncio.coroutine
    def any_to_evalscript(protocol, result):
        # Result can be native, int, MultiBulkReply or even a nested structure
        assert isinstance(result, (int, bytes, MultiBulkReply, NoneType))
        return EvalScriptReply(protocol, result)


class ListOf:
    """ Annotation helper for protocol methods. """
    def __init__(self, type_):
        self.type = type_

    def __repr__(self):
        return 'ListOf(%r)' % self.type

    def __eq__(self, other):
        return isinstance(other, ListOf) and other.type == self.type

    def __hash__(self):
        return hash((ListOf, self.type))


class NativeType:
    """
    Constant which represents the native Python type that's used.
    """
    def __new__(cls):
        raise Exception('NativeType is not meant to be initialized.')


class CommandCreator:
    """
    Utility for creating a wrapper around the Redis protocol methods.
    This will also do type checking.

    This wrapper handles (optionally) post processing of the returned data and
    implements some logic where commands behave different in case of a
    transaction or pubsub.

    Warning: We use the annotations of `method` extensively for type checking
             and determining which post processor to choose.
    """
    def __init__(self, method):
        self.method = method

    @property
    def specs(self):
       """ Argspecs """
       return getfullargspec(self.method)

    @property
    def return_type(self):
        """ Return type as defined in the method's annotation. """
        return self.specs.annotations.get('return', None)

    @property
    def params(self):
        return { k:v for k, v in self.specs.annotations.items() if k != 'return' }

    @classmethod
    def get_real_type(cls, protocol, type_):
        """
        Given a protocol instance, and type annotation, return something that
        we can pass to isinstance for the typechecking.
        """
        # If NativeType was given, replace it with the type of the protocol
        # itself.
        if isinstance(type_, tuple):
            return tuple(cls.get_real_type(protocol, t) for t in type_)

        if type_ == NativeType:
            return protocol.native_type
        elif isinstance(type_, ListOf):
            return (list, types.GeneratorType) # We don't check the content of the list.
        else:
            return type_

    def _create_input_typechecker(self):
        """ Return function that does typechecking on input data. """
        params = self.params

        if params:
            def typecheck_input(protocol, *a, **kw):
                """
                Given a protocol instance and *a/**kw of this method, raise TypeError
                when the signature doesn't match.
                """
                if protocol.enable_typechecking:
                    for name, value in getcallargs(self.method, None, *a, **kw).items():
                        if name in params:
                            real_type = self.get_real_type(protocol, params[name])
                            if not isinstance(value, real_type):
                                raise TypeError('RedisProtocol.%s received %r, expected %r' %
                                                (self.method.__name__, type(value).__name__, real_type))
        else:
            def typecheck_input(protocol, *a, **kw):
                pass

        return typecheck_input

    def _create_return_typechecker(self, return_type):
        """ Return function that does typechecking on output data. """
        if return_type and not isinstance(return_type, str): # Exclude 'Transaction'/'Subscription' which are 'str'
            def typecheck_return(protocol, result):
                """
                Given protocol and result value. Raise TypeError if the result is of the wrong type.
                """
                if protocol.enable_typechecking:
                    expected_type = self.get_real_type(protocol, return_type)
                    if not isinstance(result, expected_type):
                        raise TypeError('Got unexpected return type %r in RedisProtocol.%s, expected %r' %
                                        (type(result).__name__, self.method.__name__, expected_type))
        else:
            def typecheck_return(protocol, result):
                pass

        return typecheck_return

    def _get_docstring(self, suffix, return_type):
        # Append the real signature as the first line in the docstring.
        # (This will make the sphinx docs show the real signature instead of
        # (*a, **kw) of the wrapper.)
        # (But don't put the anotations inside the copied signature, that's rather
        # ugly in the docs.)
        signature = formatargspec(* self.specs[:6])

        # Use function annotations to generate param documentation.

        def get_name(type_):
            """ Turn type annotation into doc string. """
            try:
                return {
                    BlockingPopReply: ":class:`BlockingPopReply <asyncio_redis.replies.BlockingPopReply>`",
                    ConfigPairReply: ":class:`ConfigPairReply <asyncio_redis.replies.ConfigPairReply>`",
                    DictReply: ":class:`DictReply <asyncio_redis.replies.DictReply>`",
                    InfoReply: ":class:`InfoReply <asyncio_redis.replies.InfoReply>`",
                    ClientListReply: ":class:`InfoReply <asyncio_redis.replies.ClientListReply>`",
                    ListReply: ":class:`ListReply <asyncio_redis.replies.ListReply>`",
                    MultiBulkReply: ":class:`MultiBulkReply <asyncio_redis.replies.MultiBulkReply>`",
                    NativeType: "Native Python type, as defined by :attr:`~asyncio_redis.encoders.BaseEncoder.native_type`",
                    NoneType: "None",
                    SetReply: ":class:`SetReply <asyncio_redis.replies.SetReply>`",
                    StatusReply: ":class:`StatusReply <asyncio_redis.replies.StatusReply>`",
                    ZRangeReply: ":class:`ZRangeReply <asyncio_redis.replies.ZRangeReply>`",
                    ZScoreBoundary: ":class:`ZScoreBoundary <asyncio_redis.replies.ZScoreBoundary>`",
                    EvalScriptReply: ":class:`EvalScriptReply <asyncio_redis.replies.EvalScriptReply>`",
                    Cursor: ":class:`Cursor <asyncio_redis.cursors.Cursor>`",
                    SetCursor: ":class:`SetCursor <asyncio_redis.cursors.SetCursor>`",
                    DictCursor: ":class:`DictCursor <asyncio_redis.cursors.DictCursor>`",
                    ZCursor: ":class:`ZCursor <asyncio_redis.cursors.ZCursor>`",
                    _ScanPart: ":class:`_ScanPart",
                    int: 'int',
                    bool: 'bool',
                    dict: 'dict',
                    float: 'float',
                    str: 'str',
                    bytes: 'bytes',

                    list: 'list',
                    set: 'set',
                    dict: 'dict',

                    # XXX: Because of circulare references, we cannot use the real types here.
                    'Transaction': ":class:`asyncio_redis.Transaction`",
                    'Subscription': ":class:`asyncio_redis.Subscription`",
                    'Script': ":class:`~asyncio_redis.Script`",
                }[type_]
            except KeyError:
                if isinstance(type_, ListOf):
                    return "List or iterable of %s" % get_name(type_.type)

                elif isinstance(type_, tuple):
                    return ' or '.join(get_name(t) for t in type_)
                else:
                    raise Exception('Unknown annotation %r' % type_)
                    #return "``%s``" % type_.__name__

        def get_param(k, v):
            return ':param %s: %s\n' % (k, get_name(v))

        params_str = [ get_param(k, v) for k, v in self.params.items() ]
        returns = ':returns: (Future of) %s\n' % get_name(return_type) if return_type else ''

        return '%s%s\n%s\n\n%s%s' % (
                self.method.__name__ + suffix, signature,
                self.method.__doc__,
                ''.join(params_str),
                returns
                )

    def get_methods(self):
        """
        Return all the methods to be used in the RedisProtocol class.
        """
        return [ ('', self._get_wrapped_method(None, '', self.return_type)) ]

    def _get_wrapped_method(self, post_process, suffix, return_type):
        """
        Return the wrapped method for use in the `RedisProtocol` class.
        """
        typecheck_input = self._create_input_typechecker()
        typecheck_return = self._create_return_typechecker(return_type)
        method = self.method

        # Wrap it into a check which allows this command to be run either
        # directly on the protocol, outside of transactions or from the
        # transaction object.
        @wraps(method)
        def wrapper(protocol_self, *a, **kw):
            # When calling from a transaction
            if protocol_self.in_transaction:
                # The first arg should be the transaction # object.
                if not a or a[0] != protocol_self._transaction:
                    raise Error('Cannot run command inside transaction (use the Transaction object instead)')

                # In case of a transaction, we receive a Future from the command.
                else:
                    typecheck_input(protocol_self, *a[1:], **kw)
                    future = yield from method(protocol_self, *a[1:], **kw)
                    future2 = Future(loop=protocol_self._loop)

                    # Typecheck the future when the result is available.

                    @asyncio.coroutine
                    def done(result):
                        if post_process:
                            result = yield from post_process(protocol_self, result)
                        typecheck_return(protocol_self, result)
                        future2.set_result(result)

                    future.add_done_callback(lambda f: asyncio.async(done(f.result()), loop=protocol_self._loop))

                    return future2

            # When calling from a pubsub context
            elif protocol_self.in_pubsub:
                if not a or a[0] != protocol_self._subscription:
                    raise Error('Cannot run command inside pubsub subscription.')

                else:
                    typecheck_input(protocol_self, *a[1:], **kw)
                    result = yield from method(protocol_self, *a[1:], **kw)
                    if post_process:
                        result = yield from post_process(protocol_self, result)
                    typecheck_return(protocol_self, result)
                    return (result)

            else:
                typecheck_input(protocol_self, *a, **kw)
                result = yield from method(protocol_self, *a, **kw)
                if post_process:
                    result = yield from post_process(protocol_self, result)
                typecheck_return(protocol_self, result)
                return (result)

        wrapper.__doc__ = self._get_docstring(suffix, return_type)
        return wrapper


class QueryCommandCreator(CommandCreator):
    """
    Like `CommandCreator`, but for methods registered with `_query_command`.
    This are the methods that cause commands to be send to the server.

    Most of the commands get a reply from the server that needs to be post
    processed to get the right Python type. We inspect here the
    'returns'-annotation to determine the correct post processor.
    """
    def get_methods(self):
        # (Some commands, e.g. those that return a ListReply can generate
        # multiple protocol methods.  One that does return the ListReply, but
        # also one with the 'aslist' suffix that returns a Python list.)
        all_post_processors = PostProcessors.get_all(self.return_type)
        result = []

        for suffix, return_type, post_processor in all_post_processors:
            result.append( (suffix, self._get_wrapped_method(post_processor, suffix, return_type)) )

        return result

_SMALL_INTS = list(str(i).encode('ascii') for i in range(1000))


# List of all command methods.
_all_commands = []

class _command:
    """ Mark method as command (to be passed through CommandCreator for the
    creation of a protocol method) """
    creator = CommandCreator

    def __init__(self, method):
        self.method = method


class _query_command(_command):
    """
    Mark method as query command: This will pass through QueryCommandCreator.

    NOTE: be sure to choose the correct 'returns'-annotation. This will automatially
    determine the correct post processor function in :class:`PostProcessors`.
    """
    creator = QueryCommandCreator

    def __init__(self, method):
        super().__init__(method)


class _RedisProtocolMeta(type):
    """
    Metaclass for `RedisProtocol` which applies the _command decorator.
    """
    def __new__(cls, name, bases, attrs):
        for attr_name, value in dict(attrs).items():
            if isinstance(value, _command):
                creator = value.creator(value.method)
                for suffix, method in creator.get_methods():
                    attrs[attr_name + suffix] =  method

                    # Register command.
                    _all_commands.append(attr_name + suffix)

        return type.__new__(cls, name, bases, attrs)


class RedisProtocol(asyncio.Protocol, metaclass=_RedisProtocolMeta):
    """
    The Redis Protocol implementation.

    ::

        self.loop = asyncio.get_event_loop()
        transport, protocol = yield from loop.create_connection(RedisProtocol, 'localhost', 6379)

    :param password: <PASSWORD>
    :type password: Native Python type as defined by the ``encoder`` parameter
    :param encoder: Encoder to use for encoding to or decoding from redis bytes to a native type.
                    (Defaults to :class:`~asyncio_redis.encoders.UTF8Encoder`)
    :type encoder: :class:`~asyncio_redis.encoders.BaseEncoder` instance.
    :param db: Redis database
    :type db: int
    :param enable_typechecking: When ``True``, check argument types for all
                                redis commands. Normally you want to have this
                                enabled.
    :type enable_typechecking: bool
    """
    def __init__(self, *, password=<PASSWORD>, db=0, encoder=None, connection_lost_callback=None, enable_typechecking=True, loop=None):
        if encoder is None:
            encoder = UTF8Encoder()

        assert isinstance(db, int)
        assert isinstance(encoder, BaseEncoder)
        assert encoder.native_type, 'Encoder.native_type not defined'
        assert not password or isinstance(password, encoder.native_type)

        self.password = password
        self.db = db
        self._connection_lost_callback = connection_lost_callback
        self._loop = loop or asyncio.get_event_loop()

        # Take encode / decode settings from encoder
        self.encode_from_native = encoder.encode_from_native
        self.decode_to_native = encoder.decode_to_native
        self.native_type = encoder.native_type
        self.enable_typechecking = enable_typechecking

        self.transport = None
        self._queue = deque() # Input parser queues
        self._messages_queue = None # Pubsub queue
        self._is_connected = False # True as long as the underlying transport is connected.

        # Pubsub state
        self._in_pubsub = False
        self._subscription = None
        self._pubsub_channels = set() # Set of channels
        self._pubsub_patterns = set() # Set of patterns

        # Transaction related stuff.
        self._in_transaction = False
        self._transaction = None
        self._transaction_response_queue = None # Transaction answer queue

        self._line_received_handlers = {
            b'+': self._handle_status_reply,
            b'-': self._handle_error_reply,
            b'$': self._handle_bulk_reply,
            b'*': self._handle_multi_bulk_reply,
            b':': self._handle_int_reply,
        }

    def connection_made(self, transport):
        self.transport = transport
        self._is_connected = True
        logger.log(logging.INFO, 'Redis connection made')

        # Pipelined calls
        self._pipelined_calls = set() # Set of all the pipelined calls.

        # Start parsing reader stream.
        self._reader = StreamReader(loop=self._loop)
        self._reader.set_transport(transport)
        self._reader_f = asyncio.async(self._reader_coroutine(), loop=self._loop)

        @asyncio.coroutine
        def initialize():
            # If a password or database was been given, first connect to that one.
            if self.password:
                yield from self.auth(self.password)

            if self.db:
                yield from self.select(self.db)

            #  If we are in pubsub mode, send channel subscriptions again.
            if self._in_pubsub:
                if self._pubsub_channels:
                    yield from self._subscribe(self._subscription, list(self._pubsub_channels)) # TODO: unittest this

                if self._pubsub_patterns:
                    yield from self._psubscribe(self._subscription, list(self._pubsub_patterns))

        asyncio.async(initialize(), loop=self._loop)

    def data_received(self, data):
        """ Process data received from Redis server.  """
        self._reader.feed_data(data)

    def _encode_int(self, value:int) -> bytes:
        """ Encodes an integer to bytes. (always ascii) """
        if 0 < value < 1000: # For small values, take pre-encoded string.
            return _SMALL_INTS[value]
        else:
            return str(value).encode('ascii')

    def _encode_float(self, value:float) -> bytes:
        """ Encodes a float to bytes. (always ascii) """
        return str(value).encode('ascii')

    def _encode_zscore_boundary(self, value:ZScoreBoundary) -> str:
        """ Encodes a zscore boundary. (always ascii) """
        if isinstance(value.value, str):
            return str(value.value).encode('ascii') # +inf and -inf
        elif value.exclude_boundary:
            return str("(%f" % value.value).encode('ascii')
        else:
            return str("%f" % value.value).encode('ascii')

    def eof_received(self):
        logger.log(logging.INFO, 'EOF received in RedisProtocol')
        self._reader.feed_eof()

    def connection_lost(self, exc):
        if exc is None:
            self._reader.feed_eof()
        else:
            logger.info("Connection lost with exec: %s" % exc)
            self._reader.set_exception(exc)

        if self._reader_f:
            self._reader_f.cancel()

        self._is_connected = False
        self.transport = None
        self._reader = None
        self._reader_f = None

        # Raise exception on all waiting futures.
        while self._queue:
            f = self._queue.popleft()
            if not f.cancelled():
                f.set_exception(ConnectionLostError(exc))

        logger.log(logging.INFO, 'Redis connection lost')

        # Call connection_lost callback
        if self._connection_lost_callback:
            self._connection_lost_callback()

    # Request state

    @property
    def in_blocking_call(self):
        """ True when waiting for answer to blocking command. """
        return any(c.is_blocking for c in self._pipelined_calls)

    @property
    def in_pubsub(self):
        """ True when the protocol is in pubsub mode. """
        return self._in_pubsub

    @property
    def in_transaction(self):
        """ True when we're inside a transaction. """
        return self._in_transaction

    @property
    def in_use(self):
        """ True when this protocol is in use. """
        return self.in_blocking_call or self.in_pubsub or self.in_transaction

    @property
    def is_connected(self):
        """ True when the underlying transport is connected. """
        return self._is_connected

    # Handle replies

    @asyncio.coroutine
    def _reader_coroutine(self):
        """
        Coroutine which reads input from the stream reader and processes it.
        """
        while True:
            try:
                yield from self._handle_item(self._push_answer)
            except ConnectionLostError:
                return

    @asyncio.coroutine
    def _handle_item(self, cb):
        c = yield from self._reader.readexactly(1)
        if c:
            yield from self._line_received_handlers[c](cb)
        else:
            raise ConnectionLostError(None)

    @asyncio.coroutine
    def _handle_status_reply(self, cb):
        line = (yield from self._reader.readline()).rstrip(b'\r\n')
        cb(line)

    @asyncio.coroutine
    def _handle_int_reply(self, cb):
        line = (yield from self._reader.readline()).rstrip(b'\r\n')
        cb(int(line))

    @asyncio.coroutine
    def _handle_error_reply(self, cb):
        line = (yield from self._reader.readline()).rstrip(b'\r\n')
        cb(ErrorReply(line.decode('ascii')))

    @asyncio.coroutine
    def _handle_bulk_reply(self, cb):
        length = int((yield from self._reader.readline()).rstrip(b'\r\n'))
        if length == -1:
            # None bulk reply
            cb(None)
        else:
            # Read data
            data = yield from self._reader.readexactly(length)
            cb(data)

            # Ignore trailing newline.
            remaining = yield from self._reader.readline()
            assert remaining.rstrip(b'\r\n') == b''

    @asyncio.coroutine
    def _handle_multi_bulk_reply(self, cb):
                # NOTE: the reason for passing the callback `cb` in here is
                #       mainly because we want to return the result object
                #       especially in this case before the input is read
                #       completely. This allows a streaming API.
        count = int((yield from self._reader.readline()).rstrip(b'\r\n'))

        # Handle multi-bulk none.
        # (Used when a transaction exec fails.)
        if count == -1:
            cb(None)
            return

        reply = MultiBulkReply(self, count, loop=self._loop)

        # Return the empty queue immediately as an answer.
        if self._in_pubsub:
            asyncio.async(self._handle_pubsub_multibulk_reply(reply), loop=self._loop)
        else:
            cb(reply)

        # Wait for all multi bulk reply content.
        for i in range(count):
            yield from self._handle_item(reply._feed_received)

    @asyncio.coroutine
    def _handle_pubsub_multibulk_reply(self, multibulk_reply):
        result = yield from ListReply(multibulk_reply).aslist()
        assert result[0] in ('message', 'subscribe', 'unsubscribe', 'psubscribe', 'punsubscribe')

        if result[0] == 'message':
            channel, value = result[1], result[2]
            yield from self._subscription._messages_queue.put(PubSubReply(channel, value))

        # We can safely ignore 'subscribe'/'unsubscribe' replies at this point,
        # they don't contain anything really useful.

    # Redis operations.

    def _send_command(self, args):
        """
        Send Redis request command.
        `args` should be a list of bytes to be written to the transport.
        """
        # Create write buffer.
        data = []

        # NOTE: First, I tried to optimize by also flushing this buffer in
        # between the looping through the args. However, I removed that as the<commit_msg>Fix support for pubsub patterns.<commit_after>#!/usr/bin/env python3
import asyncio
import logging
import types

from asyncio.futures import Future
from asyncio.queues import Queue
from asyncio.streams import StreamReader

try:
    import hiredis
except ImportError:
    hiredis = None

from collections import deque
from functools import wraps
from inspect import getfullargspec, formatargspec, getcallargs

from .encoders import BaseEncoder, UTF8Encoder
from .exceptions import (
        ConnectionLostError,
        Error,
        ErrorReply,
        NoRunningScriptError,
        NotConnectedError,
        ScriptKilledError,
        TimeoutError,
        TransactionError,
)
from .log import logger
from .replies import (
        BlockingPopReply,
        ClientListReply,
        ConfigPairReply,
        DictReply,
        EvalScriptReply,
        InfoReply,
        ListReply,
        PubSubReply,
        SetReply,
        StatusReply,
        ZRangeReply,
)


from .cursors import Cursor, SetCursor, DictCursor, ZCursor

__all__ = (
    'RedisProtocol',
    'HiRedisProtocol',
    'Transaction',
    'Subscription',
    'Script',

    'ZAggregate',
    'ZScoreBoundary',
)

NoneType = type(None)


class ZScoreBoundary:
    """
    Score boundary for a sorted set.
    for queries like zrangebyscore and similar

    :param value: Value for the boundary.
    :type value: float
    :param exclude_boundary: Exclude the boundary.
    :type exclude_boundary: bool
    """
    def __init__(self, value, exclude_boundary=False):
        assert isinstance(value, float) or value in ('+inf', '-inf')
        self.value = value
        self.exclude_boundary = exclude_boundary

    def __repr__(self):
        return 'ZScoreBoundary(value=%r, exclude_boundary=%r)' % (
                    self.value, self.exclude_boundary)

ZScoreBoundary.MIN_VALUE = ZScoreBoundary('-inf')
ZScoreBoundary.MAX_VALUE = ZScoreBoundary('+inf')


class ZAggregate: # TODO: use the Python 3.4 enum type.
    """
    Aggregation method for zinterstore and zunionstore.
    """
    #: Sum aggregation.
    SUM = 'SUM'

    #: Min aggregation.
    MIN = 'MIN'

    #: Max aggregation.
    MAX = 'MAX'


class PipelinedCall:
    """ Track record for call that is being executed in a protocol. """
    __slots__ = ('cmd', 'is_blocking')

    def __init__(self, cmd, is_blocking):
        self.cmd = cmd
        self.is_blocking = is_blocking


class MultiBulkReply:
    """
    Container for a multi bulk reply.
    """
    def __init__(self, protocol, count, loop=None):
        self._loop = loop or asyncio.get_event_loop()

        #: Buffer of incoming, undelivered data, received from the parser.
        self._data_queue = []

        #: Incoming read queries.
        #: Contains (read_count, Future, decode_flag, one_only_flag) tuples.
        self._f_queue = deque()

        self.protocol = protocol
        self.count = int(count)

    def _feed_received(self, item):
        """
        Feed entry for the parser.
        """
        # Push received items on the queue
        self._data_queue.append(item)
        self._flush()

    def _flush(self):
        """
        Answer read queries when we have enough data in our multibulk reply.
        """
        # As long as we have more data in our queue then we require for a read
        # query -> answer queries.
        while self._f_queue and self._f_queue[0][0] <= len(self._data_queue):
            # Pop query.
            count, f, decode, one_only = self._f_queue.popleft()

            # Slice data buffer.
            data, self._data_queue = self._data_queue[:count], self._data_queue[count:]

            # When the decode flag is given, decode bytes to native types.
            if decode:
                data = [ self._decode(d) for d in data ]

            # When one_only flag has been given, don't return an array.
            if one_only:
                assert len(data) == 1
                f.set_result(data[0])
            else:
                f.set_result(data)

    def _decode(self, result):
        """ Decode bytes to native Python types. """
        if isinstance(result, (StatusReply, int, float, MultiBulkReply)):
            # Note that MultiBulkReplies can be nested. e.g. in the 'scan' operation.
            return result
        elif isinstance(result, bytes):
            return self.protocol.decode_to_native(result)
        elif result is None:
            return result
        else:
            raise AssertionError('Invalid type: %r' % type(result))

    def _read(self, decode=True, count=1, _one=False):
        """ Do read operation on the queue. Return future. """
        f = Future(loop=self.protocol._loop)
        self._f_queue.append((count, f, decode, _one))

        # If there is enough data on the queue, answer future immediately.
        self._flush()

        return f

    def iter_raw(self):
        """
        Iterate over all multi bulk packets. This yields futures that won't
        decode bytes yet.
        """
        for i in range(self.count):
            yield self._read(decode=False, _one=True)

    def __iter__(self):
        """
        Iterate over the reply. This yields coroutines of the decoded packets.
        It decodes bytes automatically using protocol.decode_to_native.
        """
        for i in range(self.count):
            yield self._read(_one=True)

    def __repr__(self):
        return 'MultiBulkReply(protocol=%r, count=%r)' % (self.protocol, self.count)


class _ScanPart:
    """ Internal: result chunk of a scan operation. """
    def __init__(self, new_cursor_pos, items):
        self.new_cursor_pos = new_cursor_pos
        self.items = items


class PostProcessors:
    """
    At the protocol level, we only know about a few basic classes; they
    include: bool, int, StatusReply, MultiBulkReply and bytes.
    This will return a postprocessor function that turns these into more
    meaningful objects.

    For some methods, we have several post processors. E.g. a list can be
    returned either as a ListReply (which has some special streaming
    functionality), but also as a Python list.
    """
    @classmethod
    def get_all(cls, return_type):
        """
        Return list of (suffix, return_type, post_processor)
        """
        default = cls.get_default(return_type)
        alternate = cls.get_alternate_post_processor(return_type)

        result = [ ('', return_type, default) ]
        if alternate:
            result.append(alternate)
        return result

    @classmethod
    def get_default(cls, return_type):
        """ Give post processor function for return type. """
        return {
                ListReply: cls.multibulk_as_list,
                SetReply: cls.multibulk_as_set,
                DictReply: cls.multibulk_as_dict,

                float: cls.bytes_to_float,
                (float, NoneType): cls.bytes_to_float_or_none,
                NativeType: cls.bytes_to_native,
                (NativeType, NoneType): cls.bytes_to_native_or_none,
                InfoReply: cls.bytes_to_info,
                ClientListReply: cls.bytes_to_clientlist,
                str: cls.bytes_to_str,
                bool: cls.int_to_bool,
                BlockingPopReply: cls.multibulk_as_blocking_pop_reply,
                ZRangeReply: cls.multibulk_as_zrangereply,

                StatusReply: cls.bytes_to_status_reply,
                (StatusReply, NoneType): cls.bytes_to_status_reply_or_none,
                int: None,
                (int, NoneType): None,
                ConfigPairReply: cls.multibulk_as_configpair,
                ListOf(bool): cls.multibulk_as_boolean_list,
                _ScanPart: cls.multibulk_as_scanpart,
                EvalScriptReply: cls.any_to_evalscript,

                NoneType: None,
        }[return_type]

    @classmethod
    def get_alternate_post_processor(cls, return_type):
        """ For list/set/dict. Create additional post processors that return
        python classes rather than ListReply/SetReply/DictReply """
        original_post_processor = cls.get_default(return_type)

        if return_type == ListReply:
            @asyncio.coroutine
            def as_list(protocol, result):
                result = yield from original_post_processor(protocol, result)
                return (yield from result.aslist())
            return '_aslist', list, as_list

        elif return_type == SetReply:
            @asyncio.coroutine
            def as_set(protocol, result):
                result = yield from original_post_processor(protocol, result)
                return (yield from result.asset())
            return '_asset', set, as_set

        elif return_type in (DictReply, ZRangeReply):
            @asyncio.coroutine
            def as_dict(protocol, result):
                result = yield from original_post_processor(protocol, result)
                return (yield from result.asdict())
            return '_asdict', dict, as_dict

    # === Post processor handlers below. ===

    @asyncio.coroutine
    def multibulk_as_list(protocol, result):
        assert isinstance(result, MultiBulkReply)
        return ListReply(result)

    @asyncio.coroutine
    def multibulk_as_boolean_list(protocol, result):
        # Turn the array of integers into booleans.
        assert isinstance(result, MultiBulkReply)
        values = yield from ListReply(result).aslist()
        return [ bool(v) for v in values ]

    @asyncio.coroutine
    def multibulk_as_set(protocol, result):
        assert isinstance(result, MultiBulkReply)
        return SetReply(result)

    @asyncio.coroutine
    def multibulk_as_dict(protocol, result):
        assert isinstance(result, MultiBulkReply)
        return DictReply(result)

    @asyncio.coroutine
    def multibulk_as_zrangereply(protocol, result):
        assert isinstance(result, MultiBulkReply)
        return ZRangeReply(result)

    @asyncio.coroutine
    def multibulk_as_blocking_pop_reply(protocol, result):
        if result is None:
            raise TimeoutError('Timeout in blocking pop')
        else:
            assert isinstance(result, MultiBulkReply)
            list_name, value = yield from ListReply(result).aslist()
            return BlockingPopReply(list_name, value)

    @asyncio.coroutine
    def multibulk_as_configpair(protocol, result):
        assert isinstance(result, MultiBulkReply)
        parameter, value = yield from ListReply(result).aslist()
        return ConfigPairReply(parameter, value)

    @asyncio.coroutine
    def multibulk_as_scanpart(protocol, result):
        """
        Process scanpart result.
        This is a multibulk reply of length two, where the first item is the
        new cursor position and the second item is a nested multi bulk reply
        containing all the elements.
        """
        # Get outer multi bulk reply.
        assert isinstance(result, MultiBulkReply)
        new_cursor_pos, items_bulk = yield from ListReply(result).aslist()
        assert isinstance(items_bulk, MultiBulkReply)

        # Read all items for scan chunk in memory. This is fine, because it's
        # transmitted in chunks of about 10.
        items = yield from ListReply(items_bulk).aslist()
        return _ScanPart(int(new_cursor_pos), items)

    @asyncio.coroutine
    def bytes_to_info(protocol, result):
        assert isinstance(result, bytes)
        return InfoReply(result)

    @asyncio.coroutine
    def bytes_to_status_reply(protocol, result):
        assert isinstance(result, bytes)
        return StatusReply(result.decode('utf-8'))

    @asyncio.coroutine
    def bytes_to_status_reply_or_none(protocol, result):
        assert isinstance(result, (bytes, NoneType))
        if result:
            return StatusReply(result.decode('utf-8'))

    @asyncio.coroutine
    def bytes_to_clientlist(protocol, result):
        assert isinstance(result, bytes)
        return ClientListReply(result)

    @asyncio.coroutine
    def int_to_bool(protocol, result):
        assert isinstance(result, int)
        return bool(result) # Convert int to bool

    @asyncio.coroutine
    def bytes_to_native(protocol, result):
        assert isinstance(result, bytes)
        return protocol.decode_to_native(result)

    @asyncio.coroutine
    def bytes_to_str(protocol, result):
        assert isinstance(result, bytes)
        return result.decode('ascii')

    @asyncio.coroutine
    def bytes_to_native_or_none(protocol, result):
        if result is None:
            return result
        else:
            assert isinstance(result, bytes)
            return protocol.decode_to_native(result)

    @asyncio.coroutine
    def bytes_to_float_or_none(protocol, result):
        if result is None:
            return result
        assert isinstance(result, bytes)
        return float(result)

    @asyncio.coroutine
    def bytes_to_float(protocol, result):
        assert isinstance(result, bytes)
        return float(result)

    @asyncio.coroutine
    def any_to_evalscript(protocol, result):
        # Result can be native, int, MultiBulkReply or even a nested structure
        assert isinstance(result, (int, bytes, MultiBulkReply, NoneType))
        return EvalScriptReply(protocol, result)


class ListOf:
    """ Annotation helper for protocol methods. """
    def __init__(self, type_):
        self.type = type_

    def __repr__(self):
        return 'ListOf(%r)' % self.type

    def __eq__(self, other):
        return isinstance(other, ListOf) and other.type == self.type

    def __hash__(self):
        return hash((ListOf, self.type))


class NativeType:
    """
    Constant which represents the native Python type that's used.
    """
    def __new__(cls):
        raise Exception('NativeType is not meant to be initialized.')


class CommandCreator:
    """
    Utility for creating a wrapper around the Redis protocol methods.
    This will also do type checking.

    This wrapper handles (optionally) post processing of the returned data and
    implements some logic where commands behave different in case of a
    transaction or pubsub.

    Warning: We use the annotations of `method` extensively for type checking
             and determining which post processor to choose.
    """
    def __init__(self, method):
        self.method = method

    @property
    def specs(self):
       """ Argspecs """
       return getfullargspec(self.method)

    @property
    def return_type(self):
        """ Return type as defined in the method's annotation. """
        return self.specs.annotations.get('return', None)

    @property
    def params(self):
        return { k:v for k, v in self.specs.annotations.items() if k != 'return' }

    @classmethod
    def get_real_type(cls, protocol, type_):
        """
        Given a protocol instance, and type annotation, return something that
        we can pass to isinstance for the typechecking.
        """
        # If NativeType was given, replace it with the type of the protocol
        # itself.
        if isinstance(type_, tuple):
            return tuple(cls.get_real_type(protocol, t) for t in type_)

        if type_ == NativeType:
            return protocol.native_type
        elif isinstance(type_, ListOf):
            return (list, types.GeneratorType) # We don't check the content of the list.
        else:
            return type_

    def _create_input_typechecker(self):
        """ Return function that does typechecking on input data. """
        params = self.params

        if params:
            def typecheck_input(protocol, *a, **kw):
                """
                Given a protocol instance and *a/**kw of this method, raise TypeError
                when the signature doesn't match.
                """
                if protocol.enable_typechecking:
                    for name, value in getcallargs(self.method, None, *a, **kw).items():
                        if name in params:
                            real_type = self.get_real_type(protocol, params[name])
                            if not isinstance(value, real_type):
                                raise TypeError('RedisProtocol.%s received %r, expected %r' %
                                                (self.method.__name__, type(value).__name__, real_type))
        else:
            def typecheck_input(protocol, *a, **kw):
                pass

        return typecheck_input

    def _create_return_typechecker(self, return_type):
        """ Return function that does typechecking on output data. """
        if return_type and not isinstance(return_type, str): # Exclude 'Transaction'/'Subscription' which are 'str'
            def typecheck_return(protocol, result):
                """
                Given protocol and result value. Raise TypeError if the result is of the wrong type.
                """
                if protocol.enable_typechecking:
                    expected_type = self.get_real_type(protocol, return_type)
                    if not isinstance(result, expected_type):
                        raise TypeError('Got unexpected return type %r in RedisProtocol.%s, expected %r' %
                                        (type(result).__name__, self.method.__name__, expected_type))
        else:
            def typecheck_return(protocol, result):
                pass

        return typecheck_return

    def _get_docstring(self, suffix, return_type):
        # Append the real signature as the first line in the docstring.
        # (This will make the sphinx docs show the real signature instead of
        # (*a, **kw) of the wrapper.)
        # (But don't put the anotations inside the copied signature, that's rather
        # ugly in the docs.)
        signature = formatargspec(* self.specs[:6])

        # Use function annotations to generate param documentation.

        def get_name(type_):
            """ Turn type annotation into doc string. """
            try:
                return {
                    BlockingPopReply: ":class:`BlockingPopReply <asyncio_redis.replies.BlockingPopReply>`",
                    ConfigPairReply: ":class:`ConfigPairReply <asyncio_redis.replies.ConfigPairReply>`",
                    DictReply: ":class:`DictReply <asyncio_redis.replies.DictReply>`",
                    InfoReply: ":class:`InfoReply <asyncio_redis.replies.InfoReply>`",
                    ClientListReply: ":class:`InfoReply <asyncio_redis.replies.ClientListReply>`",
                    ListReply: ":class:`ListReply <asyncio_redis.replies.ListReply>`",
                    MultiBulkReply: ":class:`MultiBulkReply <asyncio_redis.replies.MultiBulkReply>`",
                    NativeType: "Native Python type, as defined by :attr:`~asyncio_redis.encoders.BaseEncoder.native_type`",
                    NoneType: "None",
                    SetReply: ":class:`SetReply <asyncio_redis.replies.SetReply>`",
                    StatusReply: ":class:`StatusReply <asyncio_redis.replies.StatusReply>`",
                    ZRangeReply: ":class:`ZRangeReply <asyncio_redis.replies.ZRangeReply>`",
                    ZScoreBoundary: ":class:`ZScoreBoundary <asyncio_redis.replies.ZScoreBoundary>`",
                    EvalScriptReply: ":class:`EvalScriptReply <asyncio_redis.replies.EvalScriptReply>`",
                    Cursor: ":class:`Cursor <asyncio_redis.cursors.Cursor>`",
                    SetCursor: ":class:`SetCursor <asyncio_redis.cursors.SetCursor>`",
                    DictCursor: ":class:`DictCursor <asyncio_redis.cursors.DictCursor>`",
                    ZCursor: ":class:`ZCursor <asyncio_redis.cursors.ZCursor>`",
                    _ScanPart: ":class:`_ScanPart",
                    int: 'int',
                    bool: 'bool',
                    dict: 'dict',
                    float: 'float',
                    str: 'str',
                    bytes: 'bytes',

                    list: 'list',
                    set: 'set',
                    dict: 'dict',

                    # XXX: Because of circulare references, we cannot use the real types here.
                    'Transaction': ":class:`asyncio_redis.Transaction`",
                    'Subscription': ":class:`asyncio_redis.Subscription`",
                    'Script': ":class:`~asyncio_redis.Script`",
                }[type_]
            except KeyError:
                if isinstance(type_, ListOf):
                    return "List or iterable of %s" % get_name(type_.type)

                elif isinstance(type_, tuple):
                    return ' or '.join(get_name(t) for t in type_)
                else:
                    raise Exception('Unknown annotation %r' % type_)
                    #return "``%s``" % type_.__name__

        def get_param(k, v):
            return ':param %s: %s\n' % (k, get_name(v))

        params_str = [ get_param(k, v) for k, v in self.params.items() ]
        returns = ':returns: (Future of) %s\n' % get_name(return_type) if return_type else ''

        return '%s%s\n%s\n\n%s%s' % (
                self.method.__name__ + suffix, signature,
                self.method.__doc__,
                ''.join(params_str),
                returns
                )

    def get_methods(self):
        """
        Return all the methods to be used in the RedisProtocol class.
        """
        return [ ('', self._get_wrapped_method(None, '', self.return_type)) ]

    def _get_wrapped_method(self, post_process, suffix, return_type):
        """
        Return the wrapped method for use in the `RedisProtocol` class.
        """
        typecheck_input = self._create_input_typechecker()
        typecheck_return = self._create_return_typechecker(return_type)
        method = self.method

        # Wrap it into a check which allows this command to be run either
        # directly on the protocol, outside of transactions or from the
        # transaction object.
        @wraps(method)
        def wrapper(protocol_self, *a, **kw):
            # When calling from a transaction
            if protocol_self.in_transaction:
                # The first arg should be the transaction # object.
                if not a or a[0] != protocol_self._transaction:
                    raise Error('Cannot run command inside transaction (use the Transaction object instead)')

                # In case of a transaction, we receive a Future from the command.
                else:
                    typecheck_input(protocol_self, *a[1:], **kw)
                    future = yield from method(protocol_self, *a[1:], **kw)
                    future2 = Future(loop=protocol_self._loop)

                    # Typecheck the future when the result is available.

                    @asyncio.coroutine
                    def done(result):
                        if post_process:
                            result = yield from post_process(protocol_self, result)
                        typecheck_return(protocol_self, result)
                        future2.set_result(result)

                    future.add_done_callback(lambda f: asyncio.async(done(f.result()), loop=protocol_self._loop))

                    return future2

            # When calling from a pubsub context
            elif protocol_self.in_pubsub:
                if not a or a[0] != protocol_self._subscription:
                    raise Error('Cannot run command inside pubsub subscription.')

                else:
                    typecheck_input(protocol_self, *a[1:], **kw)
                    result = yield from method(protocol_self, *a[1:], **kw)
                    if post_process:
                        result = yield from post_process(protocol_self, result)
                    typecheck_return(protocol_self, result)
                    return (result)

            else:
                typecheck_input(protocol_self, *a, **kw)
                result = yield from method(protocol_self, *a, **kw)
                if post_process:
                    result = yield from post_process(protocol_self, result)
                typecheck_return(protocol_self, result)
                return (result)

        wrapper.__doc__ = self._get_docstring(suffix, return_type)
        return wrapper


class QueryCommandCreator(CommandCreator):
    """
    Like `CommandCreator`, but for methods registered with `_query_command`.
    This are the methods that cause commands to be send to the server.

    Most of the commands get a reply from the server that needs to be post
    processed to get the right Python type. We inspect here the
    'returns'-annotation to determine the correct post processor.
    """
    def get_methods(self):
        # (Some commands, e.g. those that return a ListReply can generate
        # multiple protocol methods.  One that does return the ListReply, but
        # also one with the 'aslist' suffix that returns a Python list.)
        all_post_processors = PostProcessors.get_all(self.return_type)
        result = []

        for suffix, return_type, post_processor in all_post_processors:
            result.append( (suffix, self._get_wrapped_method(post_processor, suffix, return_type)) )

        return result

_SMALL_INTS = list(str(i).encode('ascii') for i in range(1000))


# List of all command methods.
_all_commands = []

class _command:
    """ Mark method as command (to be passed through CommandCreator for the
    creation of a protocol method) """
    creator = CommandCreator

    def __init__(self, method):
        self.method = method


class _query_command(_command):
    """
    Mark method as query command: This will pass through QueryCommandCreator.

    NOTE: be sure to choose the correct 'returns'-annotation. This will automatially
    determine the correct post processor function in :class:`PostProcessors`.
    """
    creator = QueryCommandCreator

    def __init__(self, method):
        super().__init__(method)


class _RedisProtocolMeta(type):
    """
    Metaclass for `RedisProtocol` which applies the _command decorator.
    """
    def __new__(cls, name, bases, attrs):
        for attr_name, value in dict(attrs).items():
            if isinstance(value, _command):
                creator = value.creator(value.method)
                for suffix, method in creator.get_methods():
                    attrs[attr_name + suffix] =  method

                    # Register command.
                    _all_commands.append(attr_name + suffix)

        return type.__new__(cls, name, bases, attrs)


class RedisProtocol(asyncio.Protocol, metaclass=_RedisProtocolMeta):
    """
    The Redis Protocol implementation.

    ::

        self.loop = asyncio.get_event_loop()
        transport, protocol = yield from loop.create_connection(RedisProtocol, 'localhost', 6379)

    :param password: <PASSWORD>
    :type password: Native Python type as defined by the ``encoder`` parameter
    :param encoder: Encoder to use for encoding to or decoding from redis bytes to a native type.
                    (Defaults to :class:`~asyncio_redis.encoders.UTF8Encoder`)
    :type encoder: :class:`~asyncio_redis.encoders.BaseEncoder` instance.
    :param db: Redis database
    :type db: int
    :param enable_typechecking: When ``True``, check argument types for all
                                redis commands. Normally you want to have this
                                enabled.
    :type enable_typechecking: bool
    """
    def __init__(self, *, password=None, db=0, encoder=None, connection_lost_callback=None, enable_typechecking=True, loop=None):
        if encoder is None:
            encoder = UTF8Encoder()

        assert isinstance(db, int)
        assert isinstance(encoder, BaseEncoder)
        assert encoder.native_type, 'Encoder.native_type not defined'
        assert not password or isinstance(password, encoder.native_type)

        self.password = password
        self.db = db
        self._connection_lost_callback = connection_lost_callback
        self._loop = loop or asyncio.get_event_loop()

        # Take encode / decode settings from encoder
        self.encode_from_native = encoder.encode_from_native
        self.decode_to_native = encoder.decode_to_native
        self.native_type = encoder.native_type
        self.enable_typechecking = enable_typechecking

        self.transport = None
        self._queue = deque() # Input parser queues
        self._messages_queue = None # Pubsub queue
        self._is_connected = False # True as long as the underlying transport is connected.

        # Pubsub state
        self._in_pubsub = False
        self._subscription = None
        self._pubsub_channels = set() # Set of channels
        self._pubsub_patterns = set() # Set of patterns

        # Transaction related stuff.
        self._in_transaction = False
        self._transaction = None
        self._transaction_response_queue = None # Transaction answer queue

        self._line_received_handlers = {
            b'+': self._handle_status_reply,
            b'-': self._handle_error_reply,
            b'$': self._handle_bulk_reply,
            b'*': self._handle_multi_bulk_reply,
            b':': self._handle_int_reply,
        }

    def connection_made(self, transport):
        self.transport = transport
        self._is_connected = True
        logger.log(logging.INFO, 'Redis connection made')

        # Pipelined calls
        self._pipelined_calls = set() # Set of all the pipelined calls.

        # Start parsing reader stream.
        self._reader = StreamReader(loop=self._loop)
        self._reader.set_transport(transport)
        self._reader_f = asyncio.async(self._reader_coroutine(), loop=self._loop)

        @asyncio.coroutine
        def initialize():
            # If a password or database was been given, first connect to that one.
            if self.password:
                yield from self.auth(self.password)

            if self.db:
                yield from self.select(self.db)

            #  If we are in pubsub mode, send channel subscriptions again.
            if self._in_pubsub:
                if self._pubsub_channels:
                    yield from self._subscribe(self._subscription, list(self._pubsub_channels)) # TODO: unittest this

                if self._pubsub_patterns:
                    yield from self._psubscribe(self._subscription, list(self._pubsub_patterns))

        asyncio.async(initialize(), loop=self._loop)

    def data_received(self, data):
        """ Process data received from Redis server.  """
        self._reader.feed_data(data)

    def _encode_int(self, value:int) -> bytes:
        """ Encodes an integer to bytes. (always ascii) """
        if 0 < value < 1000: # For small values, take pre-encoded string.
            return _SMALL_INTS[value]
        else:
            return str(value).encode('ascii')

    def _encode_float(self, value:float) -> bytes:
        """ Encodes a float to bytes. (always ascii) """
        return str(value).encode('ascii')

    def _encode_zscore_boundary(self, value:ZScoreBoundary) -> str:
        """ Encodes a zscore boundary. (always ascii) """
        if isinstance(value.value, str):
            return str(value.value).encode('ascii') # +inf and -inf
        elif value.exclude_boundary:
            return str("(%f" % value.value).encode('ascii')
        else:
            return str("%f" % value.value).encode('ascii')

    def eof_received(self):
        logger.log(logging.INFO, 'EOF received in RedisProtocol')
        self._reader.feed_eof()

    def connection_lost(self, exc):
        if exc is None:
            self._reader.feed_eof()
        else:
            logger.info("Connection lost with exec: %s" % exc)
            self._reader.set_exception(exc)

        if self._reader_f:
            self._reader_f.cancel()

        self._is_connected = False
        self.transport = None
        self._reader = None
        self._reader_f = None

        # Raise exception on all waiting futures.
        while self._queue:
            f = self._queue.popleft()
            if not f.cancelled():
                f.set_exception(ConnectionLostError(exc))

        logger.log(logging.INFO, 'Redis connection lost')

        # Call connection_lost callback
        if self._connection_lost_callback:
            self._connection_lost_callback()

    # Request state

    @property
    def in_blocking_call(self):
        """ True when waiting for answer to blocking command. """
        return any(c.is_blocking for c in self._pipelined_calls)

    @property
    def in_pubsub(self):
        """ True when the protocol is in pubsub mode. """
        return self._in_pubsub

    @property
    def in_transaction(self):
        """ True when we're inside a transaction. """
        return self._in_transaction

    @property
    def in_use(self):
        """ True when this protocol is in use. """
        return self.in_blocking_call or self.in_pubsub or self.in_transaction

    @property
    def is_connected(self):
        """ True when the underlying transport is connected. """
        return self._is_connected

    # Handle replies

    @asyncio.coroutine
    def _reader_coroutine(self):
        """
        Coroutine which reads input from the stream reader and processes it.
        """
        while True:
            try:
                yield from self._handle_item(self._push_answer)
            except ConnectionLostError:
                return

    @asyncio.coroutine
    def _handle_item(self, cb):
        c = yield from self._reader.readexactly(1)
        if c:
            yield from self._line_received_handlers[c](cb)
        else:
            raise ConnectionLostError(None)

    @asyncio.coroutine
    def _handle_status_reply(self, cb):
        line = (yield from self._reader.readline()).rstrip(b'\r\n')
        cb(line)

    @asyncio.coroutine
    def _handle_int_reply(self, cb):
        line = (yield from self._reader.readline()).rstrip(b'\r\n')
        cb(int(line))

    @asyncio.coroutine
    def _handle_error_reply(self, cb):
        line = (yield from self._reader.readline()).rstrip(b'\r\n')
        cb(ErrorReply(line.decode('ascii')))

    @asyncio.coroutine
    def _handle_bulk_reply(self, cb):
        length = int((yield from self._reader.readline()).rstrip(b'\r\n'))
        if length == -1:
            # None bulk reply
            cb(None)
        else:
            # Read data
            data = yield from self._reader.readexactly(length)
            cb(data)

            # Ignore trailing newline.
            remaining = yield from self._reader.readline()
            assert remaining.rstrip(b'\r\n') == b''

    @asyncio.coroutine
    def _handle_multi_bulk_reply(self, cb):
                # NOTE: the reason for passing the callback `cb` in here is
                #       mainly because we want to return the result object
                #       especially in this case before the input is read
                #       completely. This allows a streaming API.
        count = int((yield from self._reader.readline()).rstrip(b'\r\n'))

        # Handle multi-bulk none.
        # (Used when a transaction exec fails.)
        if count == -1:
            cb(None)
            return

        reply = MultiBulkReply(self, count, loop=self._loop)

        # Return the empty queue immediately as an answer.
        if self._in_pubsub:
            asyncio.async(self._handle_pubsub_multibulk_reply(reply), loop=self._loop)
        else:
            cb(reply)

        # Wait for all multi bulk reply content.
        for i in range(count):
            yield from self._handle_item(reply._feed_received)

    @asyncio.coroutine
    def _handle_pubsub_multibulk_reply(self, multibulk_reply):
        result = yield from ListReply(multibulk_reply).aslist()
        assert result[0] in ('message', 'subscribe', 'unsubscribe', 'pmessage', 'psubscribe', 'punsubscribe')

        if result[0] == 'message':
            channel, value = result[1], result[2]
            yield from self._subscription._messages_queue.put(PubSubReply(channel, value))
        elif result[0] == 'pmessage':
            channel, value = result[2], result[3]
            yield from self._subscription._messages_queue.put(PubSubReply(channel, value))

        # We can safely ignore 'subscribe'/'unsubscribe' replies at this point,
        # they don't contain anything really useful.

    # Redis operations.

    def _send_command(self, args):
        """
        Send Redis request command.
        `args` should be a list of bytes to be written to the transport.
        """
        # Create write buffer.
        data = []

        # NOTE: First, I tried to optimize by also flushing this buffer in
        # between the looping through the args. However, I removed that as the
{'fraction_non_alphanumeric': 0.08726132024004364, 'fraction_numerical': 0.0012956901254773596, 'mean_word_length': 2.1938406586226424, 'pattern_counts': {'":': 40, '<': 44, '<?xml version=': 0, '>': 46, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1428528', 'n_tokens_mistral': 20298, 'n_tokens_neox': 19446, 'n_words': 5707}