<commit_before>import asyncio

import psycopg2
from psycopg2.extensions import (
    POLL_OK, POLL_READ, POLL_WRITE, POLL_ERROR)
from .exceptions import UnknownPollError, ConnectionClosedError


__all__ = ('connect',)


ALLOWED_ARGS = {'host', 'hostaddr', 'port', 'dbname', 'user',
                'password', 'connect_timeout', 'client_encoding',
                'options', 'application_name',
                'fallback_application_name', 'keepalives',
                'keepalives_idle', 'keepalives_interval',
                'keepalives_count', 'tty', 'sslmode', 'requiressl',
                'sslcompression', 'sslcert', 'sslkey', 'sslrootcert',
                'sslcrl', 'requirepeer', 'krbsrvname', 'gsslib',
                'service', 'database', 'connection_factory', 'cursor_factory'}


@asyncio.coroutine
def connect(dsn=None, *,
            loop=None, **kwargs):
    """XXX"""

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

    for k in kwargs:
        if k not in ALLOWED_ARGS:
            raise TypeError("connect() got unexpected keyword argument '{}'"
                            .format(k))

    waiter = asyncio.Future(loop=loop)
    conn = Connection(dsn, loop, waiter, **kwargs)
    yield from waiter
    return conn


class Connection:
    """Psycopg connection wrapper class.
    :param string dsn:
    :param psycopg2.extensions.cursor cursor_factory: argument can be used
        to create non-standard cursors
    :param psycopg2.extensions.connection connection_factory: class is usually
        sub-classed only to provide an easy way to create customized cursors
        but other uses are possible
    :param asyncio.EventLoop loop: A list or tuple with query parameters.
            Defaults to an empty tuple."""
    def __init__(self, dsn, loop, waiter,
                 **kwargs):

        self._loop = loop

        self._conn = psycopg2.connect(
            dsn,
            async=True,
            **kwargs)
        self._fileno = self._conn.fileno()
        self._waiter = waiter
        self._ready(None)

    def _ready(self, action):
        assert self._waiter is not None, "BAD STATE"
        print('READY', action)

        if action == None:
            pass
        elif action == 'writing':
            self._loop.remove_writer(self._fileno)
        elif action == 'reading':
            self._loop.remove_reader(self._fileno)
        else:
            self._fatal_error(RuntimeError("Unknown action {!r}"
                                           .format(action)))

        try:
            state = self._conn.poll()
            print('READY STATE', state)
        except (psycopg2.Warning, psycopg2.Error) as error:
            self._waiter.set_exception(error)
            self._waiter = None
        else:
            if state == POLL_OK:
                print('READY DONE')
                self._waiter.set_result(None)
                self._waiter = None
            elif state == POLL_READ:
                print('READY READ')<commit_msg>Workin on<commit_after>import asyncio

import psycopg2
from psycopg2.extensions import (
    POLL_OK, POLL_READ, POLL_WRITE, POLL_ERROR)
from .exceptions import UnknownPollError, ConnectionClosedError


__all__ = ('connect',)


ALLOWED_ARGS = {'host', 'hostaddr', 'port', 'dbname', 'user',
                'password', 'connect_timeout', 'client_encoding',
                'options', 'application_name',
                'fallback_application_name', 'keepalives',
                'keepalives_idle', 'keepalives_interval',
                'keepalives_count', 'tty', 'sslmode', 'requiressl',
                'sslcompression', 'sslcert', 'sslkey', 'sslrootcert',
                'sslcrl', 'requirepeer', 'krbsrvname', 'gsslib',
                'service', 'database', 'connection_factory', 'cursor_factory'}


#NB: psycopg2 adds POLL_READ and POLL_WRITE to self._conn.poll() and never rejects the options until POLL_OK. We need to process that.


@asyncio.coroutine
def connect(dsn=None, *, loop=None, **kwargs):
    """XXX"""

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

    for k in kwargs:
        if k not in ALLOWED_ARGS:
            raise TypeError("connect() got unexpected keyword argument '{}'"
                            .format(k))

    waiter = asyncio.Future(loop=loop)
    conn = Connection(dsn, loop, waiter, **kwargs)
    yield from waiter
    return conn


class Connection:
    """XXX"""

    def __init__(self, dsn, loop, waiter,
                 **kwargs):

        self._loop = loop

        self._conn = psycopg2.connect(
            dsn,
            async=True,
            **kwargs)
        self._fileno = self._conn.fileno()
        self._waiter = waiter
        self._ready('writing')

    def _ready(self, action):
        assert self._waiter is not None, "BAD STATE"
        print('READY', action)

        if action == None:
            pass
        elif action == 'writing':
            self._loop.remove_writer(self._fileno)
        elif action == 'reading':
            self._loop.remove_reader(self._fileno)
        else:
            self._fatal_error(RuntimeError("Unknown action {!r}"
                                           .format(action)))

        try:
            state = self._conn.poll()
            print('READY STATE', state)
        except (psycopg2.Warning, psycopg2.Error) as error:
            self._waiter.set_exception(error)
            self._waiter = None
        else:
            if state == POLL_OK:
                if not self._conn.isexecuted():
                    print('READY DONE')
                    self._waiter.set_result(None)
                    self._waiter = None
            elif state == POLL_READ:
                print('READY READ')
{'fraction_non_alphanumeric': 0.10841480446927375, 'fraction_numerical': 0.0022695530726256985, 'mean_word_length': 2.012092534174553, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2053138', 'n_tokens_mistral': 1686, 'n_tokens_neox': 1596, 'n_words': 411}