<commit_before>import types
import asyncio
import socket
from functools import partial
from collections import deque

from .util import (
    encode_command,
    wait_ok,
    _NOTSET,
    _set_result,
    _set_exception,
    coerced_keys_dict,
    decode,
    async_task,
    create_future,
    )
from .parser import Reader
from .stream import open_connection, open_unix_connection
from .errors import (
    ConnectionClosedError,
    RedisError,
    ProtocolError,
    ReplyError,
    WatchVariableError,
    ReadOnlyError,
    )
from .pubsub import Channel
from .abc import AbcChannel
from .abc import AbcConnection
from .log import logger


__all__ = ['create_connection', 'RedisConnection']

MAX_CHUNK_SIZE = 65536

_PUBSUB_COMMANDS = (
    'SUBSCRIBE', b'SUBSCRIBE',
    'PSUBSCRIBE', b'PSUBSCRIBE',
    'UNSUBSCRIBE', b'UNSUBSCRIBE',
    'PUNSUBSCRIBE', b'PUNSUBSCRIBE',
    )


@asyncio.coroutine
def create_connection(address, *, db=None, password=<PASSWORD>, ssl=None,
                      encoding=None, parser=None, loop=None, timeout=None,
                      connection_cls=None):
    """Creates redis connection.

    Opens connection to Redis server specified by address argument.
    Address argument is similar to socket address argument, ie:
    * when address is a tuple it represents (host, port) pair;
    * when address is a str it represents unix domain socket path.
    (no other address formats supported)

    SSL argument is passed through to asyncio.create_connection.
    By default SSL/TLS is not used.

    By default any timeout is applied at the connection stage, however
    you can set a limitted time used trying to open a connection via
    the `timeout` Kw.

    Encoding argument can be used to decode byte-replies to strings.
    By default no decoding is done.

    Parser parameter can be used to pass custom Redis protocol parser class.
    By default hiredis.Reader is used (unless it is missing or platform
    is not CPython).

    Return value is RedisConnection instance or a connection_cls if it is
    given.

    This function is a coroutine.
    """
    assert isinstance(address, (tuple, list, str)), "tuple or str expected"

    if timeout is not None and timeout <= 0:
        raise ValueError("Timeout has to be None or a number greater than 0")

    if connection_cls:
        assert issubclass(connection_cls, AbcConnection),\
                "connection_class does not meet the AbcConnection contract"
        cls = connection_cls
    else:
        cls = RedisConnection

    if loop is None:
        loop = asyncio.get_event_loop()

    if isinstance(address, (list, tuple)):
        host, port = address
        logger.debug("Creating tcp connection to %r", address)
        reader, writer = yield from asyncio.wait_for(open_connection(
            host, port, limit=MAX_CHUNK_SIZE, ssl=ssl, loop=loop),
            timeout, loop=loop)
        sock = writer.transport.get_extra_info('socket')
        if sock is not None:
            sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
            address = sock.getpeername()
        address = tuple(address[:2])
    else:
        logger.debug("Creating unix connection to %r", address)
        reader, writer = yield from asyncio.wait_for(open_unix_connection(
            address, ssl=ssl, limit=MAX_CHUNK_SIZE, loop=loop),
            timeout, loop=loop)
        sock = writer.transport.get_extra_info('socket')
        if sock is not None:
            address = sock.getpeername()

    conn = cls(reader, writer, encoding=encoding,
               address=address, parser=parser,
               loop=loop)

    try:
        if password is not None:
            yield from conn.auth(password)
        if db is not None:
            yield from conn.select(db)
    except Exception:
        conn.close()
        yield from conn.wait_closed()
        raise
    return conn


class RedisConnection(AbcConnection):
    """Redis connection."""

    def __init__(self, reader, writer, *, address, encoding=None,
                 parser=None, loop=None):
        if loop is None:
            loop = asyncio.get_event_loop()
        if parser is None:
            parser = Reader
        assert callable(parser), (
            "Parser argument is not callable", parser)
        self._reader = reader
        self._writer = writer
        self._address = address
        self._loop = loop
        self._waiters = deque()
        self._reader.set_parser(
            parser(protocolError=ProtocolError, replyError=ReplyError)
        )
        self._reader_task = async_task(self._read_data(), loop=self._loop)
        self._db = 0
        self._closing = False
        self._closed = False
        self._close_waiter = create_future(loop=self._loop)
        self._reader_task.add_done_callback(self._close_waiter.set_result)
        self._in_transaction = None
        self._transaction_error = None  # XXX: never used?
        self._in_pubsub = 0
        self._pubsub_channels = coerced_keys_dict()
        self._pubsub_patterns = coerced_keys_dict()
        self._encoding = encoding

    def __repr__(self):
        return '<RedisConnection [db:{}]>'.format(self._db)

    @asyncio.coroutine
    def _read_data(self):
        """Response reader task."""
        while not self._reader.at_eof():
            try:
                obj = yield from self._reader.readobj()
            except asyncio.CancelledError:
                break
            except ProtocolError as exc:
                # ProtocolError is fatal
                # so connection must be closed
                if self._in_transaction is not None:
                    self._transaction_error = exc
                self._closing = True
                self._do_close(exc)
                return
            except Exception as exc:
                # XXX: for QUIT command connection error can be received
                #       before response
                logger.error("Exception on data read %r", exc, exc_info=True)
                break
            else:
                if (obj == b'' or obj is None) and self._reader.at_eof():
                    logger.debug("Connection has been closed by server,"
                                 " response: %r", obj)
                    break

                if self._in_pubsub:
                    self._process_pubsub(obj)
                else:
                    self._process_data(obj)

        self._closing = True
        self._do_close(None)

    def _process_data(self, obj):
        """Processes command results."""
        assert len(self._waiters) > 0, (type(obj), obj)
        waiter, encoding, cb = self._waiters.popleft()
        if isinstance(obj, RedisError):
            if isinstance(obj, ReplyError):
                if obj.args[0].startswith('READONLY'):
                    obj = ReadOnlyError(obj.args[0])
            _set_exception(waiter, obj)
            if self._in_transaction is not None:
                self._transaction_error = obj
        else:
            if encoding is not None:
                try:
                    obj = decode(obj, encoding)
                except Exception as exc:
                    _set_exception(waiter, exc)
                    return<commit_msg>Use call_soon() when closing in conn._read_data()<commit_after>import types
import asyncio
import socket
from functools import partial
from collections import deque

from .util import (
    encode_command,
    wait_ok,
    _NOTSET,
    _set_result,
    _set_exception,
    coerced_keys_dict,
    decode,
    async_task,
    create_future,
    )
from .parser import Reader
from .stream import open_connection, open_unix_connection
from .errors import (
    ConnectionClosedError,
    RedisError,
    ProtocolError,
    ReplyError,
    WatchVariableError,
    ReadOnlyError,
    )
from .pubsub import Channel
from .abc import AbcChannel
from .abc import AbcConnection
from .log import logger


__all__ = ['create_connection', 'RedisConnection']

MAX_CHUNK_SIZE = 65536

_PUBSUB_COMMANDS = (
    'SUBSCRIBE', b'SUBSCRIBE',
    'PSUBSCRIBE', b'PSUBSCRIBE',
    'UNSUBSCRIBE', b'UNSUBSCRIBE',
    'PUNSUBSCRIBE', b'PUNSUBSCRIBE',
    )


@asyncio.coroutine
def create_connection(address, *, db=None, password=<PASSWORD>, ssl=None,
                      encoding=None, parser=None, loop=None, timeout=None,
                      connection_cls=None):
    """Creates redis connection.

    Opens connection to Redis server specified by address argument.
    Address argument is similar to socket address argument, ie:
    * when address is a tuple it represents (host, port) pair;
    * when address is a str it represents unix domain socket path.
    (no other address formats supported)

    SSL argument is passed through to asyncio.create_connection.
    By default SSL/TLS is not used.

    By default any timeout is applied at the connection stage, however
    you can set a limitted time used trying to open a connection via
    the `timeout` Kw.

    Encoding argument can be used to decode byte-replies to strings.
    By default no decoding is done.

    Parser parameter can be used to pass custom Redis protocol parser class.
    By default hiredis.Reader is used (unless it is missing or platform
    is not CPython).

    Return value is RedisConnection instance or a connection_cls if it is
    given.

    This function is a coroutine.
    """
    assert isinstance(address, (tuple, list, str)), "tuple or str expected"

    if timeout is not None and timeout <= 0:
        raise ValueError("Timeout has to be None or a number greater than 0")

    if connection_cls:
        assert issubclass(connection_cls, AbcConnection),\
                "connection_class does not meet the AbcConnection contract"
        cls = connection_cls
    else:
        cls = RedisConnection

    if loop is None:
        loop = asyncio.get_event_loop()

    if isinstance(address, (list, tuple)):
        host, port = address
        logger.debug("Creating tcp connection to %r", address)
        reader, writer = yield from asyncio.wait_for(open_connection(
            host, port, limit=MAX_CHUNK_SIZE, ssl=ssl, loop=loop),
            timeout, loop=loop)
        sock = writer.transport.get_extra_info('socket')
        if sock is not None:
            sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
            address = sock.getpeername()
        address = tuple(address[:2])
    else:
        logger.debug("Creating unix connection to %r", address)
        reader, writer = yield from asyncio.wait_for(open_unix_connection(
            address, ssl=ssl, limit=MAX_CHUNK_SIZE, loop=loop),
            timeout, loop=loop)
        sock = writer.transport.get_extra_info('socket')
        if sock is not None:
            address = sock.getpeername()

    conn = cls(reader, writer, encoding=encoding,
               address=address, parser=parser,
               loop=loop)

    try:
        if password is not None:
            yield from conn.auth(password)
        if db is not None:
            yield from conn.select(db)
    except Exception:
        conn.close()
        yield from conn.wait_closed()
        raise
    return conn


class RedisConnection(AbcConnection):
    """Redis connection."""

    def __init__(self, reader, writer, *, address, encoding=None,
                 parser=None, loop=None):
        if loop is None:
            loop = asyncio.get_event_loop()
        if parser is None:
            parser = Reader
        assert callable(parser), (
            "Parser argument is not callable", parser)
        self._reader = reader
        self._writer = writer
        self._address = address
        self._loop = loop
        self._waiters = deque()
        self._reader.set_parser(
            parser(protocolError=ProtocolError, replyError=ReplyError)
        )
        self._reader_task = async_task(self._read_data(), loop=self._loop)
        self._db = 0
        self._closing = False
        self._closed = False
        self._close_waiter = create_future(loop=self._loop)
        self._reader_task.add_done_callback(self._close_waiter.set_result)
        self._in_transaction = None
        self._transaction_error = None  # XXX: never used?
        self._in_pubsub = 0
        self._pubsub_channels = coerced_keys_dict()
        self._pubsub_patterns = coerced_keys_dict()
        self._encoding = encoding

    def __repr__(self):
        return '<RedisConnection [db:{}]>'.format(self._db)

    @asyncio.coroutine
    def _read_data(self):
        """Response reader task."""
        while not self._reader.at_eof():
            try:
                obj = yield from self._reader.readobj()
            except asyncio.CancelledError:
                break
            except ProtocolError as exc:
                # ProtocolError is fatal
                # so connection must be closed
                if self._in_transaction is not None:
                    self._transaction_error = exc
                self._closing = True
                self._loop.call_soon(self._do_close, exc)
                return
            except Exception as exc:
                # XXX: for QUIT command connection error can be received
                #       before response
                logger.error("Exception on data read %r", exc, exc_info=True)
                break
            else:
                if (obj == b'' or obj is None) and self._reader.at_eof():
                    logger.debug("Connection has been closed by server,"
                                 " response: %r", obj)
                    break

                if self._in_pubsub:
                    self._process_pubsub(obj)
                else:
                    self._process_data(obj)

        self._closing = True
        self._loop.call_soon(self._do_close, None)

    def _process_data(self, obj):
        """Processes command results."""
        assert len(self._waiters) > 0, (type(obj), obj)
        waiter, encoding, cb = self._waiters.popleft()
        if isinstance(obj, RedisError):
            if isinstance(obj, ReplyError):
                if obj.args[0].startswith('READONLY'):
                    obj = ReadOnlyError(obj.args[0])
            _set_exception(waiter, obj)
            if self._in_transaction is not None:
                self._transaction_error = obj
        else:
            if encoding is not None:
                try:
                    obj = decode(obj, encoding)
                except Exception as exc:
                    _set_exception(waiter, exc)
                    return
{'fraction_non_alphanumeric': 0.07805849237951394, 'fraction_numerical': 0.0019222847727584787, 'mean_word_length': 2.1039846580012784, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4000095', 'n_tokens_mistral': 4051, 'n_tokens_neox': 3905, 'n_words': 1208}