<commit_before>"""The core module is transfered to the remote process and will bootstrap pipe communication.

It creates default channels and dispatches commands accordingly.

"""  # pylint: disable=C0302
import asyncio
import collections
import concurrent
import functools
import hashlib
import importlib.abc
import importlib.machinery
import importlib.util
import logging
import logging.config
import os
import signal
import struct
import sys
import time
import traceback
import uuid
import weakref
import zlib

from . import msgpack

PY_35 = sys.version_info >= (3, 5)
PY_352 = sys.version_info >= (3, 5, 2)

log = logging.getLogger(__name__)


@msgpack.register(object, 0x01)
class CustomEncoder:
    @classmethod
    def __msgpack_encode__(cls, data, data_type):
        data_type = type(data)
        encoder = msgpack.get_custom_encoder(data_type)
        if encoder is None:
            raise TypeError("There is no custom encoder for this type registered: {}".format(data_type))

        wrapped = {
            'type': data_type.__name__,
            'module': data_type.__module__,
            'data': encoder.__msgpack_encode__(data, data_type)
        }
        return msgpack.encode(wrapped)

    @classmethod
    def __msgpack_decode__(cls, encoded_data, data_type):
        wrapped = msgpack.decode(encoded_data)

        try:
            module = sys.modules[wrapped['module']]
        except KeyError:
            raise
            # XXX should we import the module?

        data_type = getattr(module, wrapped['type'])
        encoder = msgpack.get_custom_encoder(data_type)
        if encoder is None:
            raise TypeError("There is no custom encoder for this type registered: {}".format(data_type))

        data = encoder.__msgpack_decode__(wrapped['data'], data_type)
        return data


@msgpack.register(tuple, 0x02)
class TupleEncoder:
    @classmethod
    def __msgpack_encode__(cls, data, data_type):
        return msgpack.encode(list(data))

    @classmethod
    def __msgpack_decode__(cls, encoded_data, data_type):
        return tuple(msgpack.decode(encoded_data))


@msgpack.register(set, 0x03)
class SetEncoder:
    @classmethod
    def __msgpack_encode__(cls, data, data_type):
        return msgpack.encode(list(data))

    @classmethod
    def __msgpack_decode__(cls, encoded_data, data_type):
        return set(msgpack.decode(encoded_data))


@msgpack.register(Exception, 0x04)
class ExceptionEncoder:
    @classmethod
    def __msgpack_encode__(cls, data, data_type):
        return msgpack.encode(data.args)

    @classmethod
    def __msgpack_decode__(cls, encoded_data, data_type):
        return data_type(*msgpack.decode(encoded_data))


@msgpack.register(StopAsyncIteration, 0x05)
class StopAsyncIterationEncoder:
    @classmethod
    def __msgpack_encode__(cls, data, data_type):
        return msgpack.encode(data.args)

    @classmethod
    def __msgpack_decode__(cls, encoded_data, data_type):
        return StopAsyncIteration(*msgpack.decode(encoded_data))


class Uid(uuid.UUID):
    def __init__(self, bytes=None):
        if bytes is None:
            super().__init__(bytes=uuid.uuid1().bytes, version=1)
        else:
            super().__init__(bytes=bytes, version=1)

    @property
    def time(self):
        time = (super().time - 0x01b21dd213814000) * 100 / 1e9
        return time

    @classmethod
    def __msgpack_encode__(cls, data, data_type):
        return data.bytes

    @classmethod
    def __msgpack_decode__(cls, encoded, data_type):
        return cls(bytes=encoded)


class ConnectionLostStreamReaderProtocol(asyncio.StreamReaderProtocol):
    def __init__(self, *args, connection_lost_cb, **kwargs):
        super().__init__(*args, **kwargs)
        self.connection_lost_cb = connection_lost_cb

    def connection_lost(self, exc):
        super().connection_lost(exc)
        self.connection_lost_cb(exc)


class Incomming(asyncio.StreamReader):

    """A context for an incomming pipe."""

    def __init__(self, *, connection_lost_cb=None, pipe=sys.stdin, loop=None):
        super(Incomming, self).__init__(loop=loop)
        self.pipe = os.fdopen(pipe) if isinstance(pipe, int) else pipe
        self.connection_lost_cb = connection_lost_cb

    async def __aenter__(self):
        await self.connect()
        return self

    async def connect(self):
        if self.connection_lost_cb:
            protocol = ConnectionLostStreamReaderProtocol(
                self, connection_lost_cb=self.connection_lost_cb, loop=self._loop
            )
        else:
            protocol = asyncio.StreamReaderProtocol(self, loop=self._loop)

        transport, protocol = await self._loop.connect_read_pipe(
            lambda: protocol,
            self.pipe,
        )

        return transport, protocol

    async def __aexit__(self, exc_type, value, tb):
        self._transport.close()

    async def readexactly(self, n):
        """Read exactly n bytes from the stream.

        This is a short and faster implementation then the original one
        (see of https://github.com/python/asyncio/issues/394).

        """
        buffer, missing = bytearray(), n

        while missing:
            if not self._buffer:
                await self._wait_for_data('readexactly')

            if self._eof or not self._buffer:
                raise asyncio.IncompleteReadError(bytes(buffer), n)

            length = min(len(self._buffer), missing)
            buffer.extend(self._buffer[:length])

            del self._buffer[:length]
            missing -= length

            self._maybe_resume_transport()

        return buffer


class Outgoing:

    """A context for an outgoing pipe."""

    def __init__(self, *, pipe=sys.stdout, reader=None, loop=None):
        self.loop = loop if loop is not None else asyncio.get_event_loop()
        self.pipe = os.fdopen(pipe) if isinstance(pipe, int) else pipe
        self.transport = None
        self.reader = reader
        self.writer = None

    async def __aenter__(self):
        writer = await self.connect()
        return writer

    async def connect(self):
        self.transport, protocol = await self.loop.connect_write_pipe(
            asyncio.streams.FlowControlMixin,
            self.pipe
        )

        writer = asyncio.streams.StreamWriter(
            self.transport,
            protocol,
            self.reader,
            self.loop
        )
        return writer

    async def __aexit__(self, exc_type, value, tb):
        self.transport.close()


def split_data(data, size=1024):
    """Create a generator to split data into chunks."""
    data_view, data_len, start = memoryview(data), len(data), 0
    while start < data_len:
        end = min(start + size, data_len)
        yield data_view[start:end]
        start = end


class OrderedMeta(type):
    items = []

    def __call__(cls, *args, **kwargs):
        inst = super().__call__(*args, **kwargs)
        cls.items.append(inst)
        return inst

    @classmethod
    def ordered_items(mcs, dct, cls_order=None):
        def sort(item):
            if cls_order is not None:
                try:
                    cls_index = cls_order.index(item[1].__class__)
                except ValueError:
                    cls_index = None
                return cls_index, mcs.items.index(item[1])
            return mcs.items.index(item[1])

        items = collections.OrderedDict(sorted(
            (item for item in dct.items() if item[1] in mcs.items),
            key=sort
        ))
        return items


class BaseHeaderItem(metaclass=OrderedMeta):
    def __init__(self, *, default=None, encoder=None, decoder=None):
        self.index = None
        self.name = None
        self.default = default
        self.encoder = encoder
        self.decoder = decoder

    def set_index(self, index):
        self.index = index

    def __set_name__(self, owner, name):
        self.name = name

    def encode(self, value):
        if self.encoder:
            value = self.encoder(value, self)
        return value

    def decode(self, value):
        return value

    def __get__(self, inst, cls):
        if inst is None:
            return self
        value = self.decode(inst)
        if self.decoder:
            value = self.decoder(value, self)
        setattr(inst, self.name, value)
        return value


class HeaderItem(BaseHeaderItem):
    def __init__(self, fmt, **kwargs):
        super().__init__(**kwargs)
        self.fmt = fmt

    @property
    def size(self):
        return struct.calcsize(self.fmt)

    def encode(self, value):
        value = super().encode(value)
        return struct.pack(self.fmt, value or self.default)

    def decode(self, value):
        value, = struct.unpack_from(self.fmt, value, self.index)
        return value


class Flag(BaseHeaderItem):
    def encode(self, value):
        value = super().encode(value)
        return 0 if not value else value << self.index

    def decode(self, value):
        value, = struct.unpack_from('!1B', value, 0)
        value = value >> self.index & 1
        return value


class HeaderMeta(type):
    def __new__(mcs, name, bases, dct):
        items = dct['items'] = OrderedMeta.ordered_items(dct, (Flag, HeaderItem))
        dct['size'] = mcs.apply_items_index(items)
        cls = type.__new__(mcs, name, bases, dct)
        if sys.version_info < (3, 6):
            for name, item in items.items():
                item.__set_name__(cls, name)
        return cls

    @classmethod
    def apply_items_index(mcs, items):
        flag_count = 0
        item_size = 0
        for name, item in items.items():
            if isinstance(item, Flag):
                item.set_index(flag_count)
                flag_count += 1
                item_size = (flag_count // 8) + (flag_count % 8 != 0)
            else:
                # items are ordered after type
                item.set_index(item_size)
                item_size += item.size
        return item_size

    def __call__(cls, encoded_bytes=None, **items):
        if encoded_bytes:
            inst = super().__call__(encoded_bytes)
            return inst

        flags = 0
        encoded_items = []
        for name, item in cls.items.items():
            if isinstance(item, Flag):
                flags += item.encode(items.get(name, False))
            else:
                encoded_items.append(item.encode(items.get(name, None)))
        encoded_flags = bytes([flags])
        encoded_items.insert(0, encoded_flags)

        inst = super().__call__(b''.join(encoded_items))
        return inst


class Header(bytes, metaclass=HeaderMeta):
    eom = Flag()
    send_ack = Flag()
    recv_ack = Flag()
    compression = Flag()
    uid = HeaderItem('!16s', encoder=lambda uid, _: uid.bytes, decoder=lambda x, _: Uid(bytes=x))
    channel_name_len = HeaderItem('!H', default=0)
    data_len = HeaderItem('!I', default=0)

    def __repr__(self):
        return '<Header {self.uid} {eom} {self.data_len}>'.format(self=self, eom=' eom' if self.eom else '')


Chunk = collections.namedtuple('Chunk', ('header', 'channel_name', 'data'))


class Channels:

    """Hold references to all channel queues."""<commit_msg>docs and cleanup<commit_after>"""The core module is transfered to the remote process and will bootstrap pipe communication.

It creates default channels and dispatches commands accordingly.

"""  # pylint: disable=C0302
import asyncio
import collections
import concurrent
import functools
import importlib.abc
import importlib.machinery
import importlib.util
import logging
import logging.config
import os
import signal
import struct
import sys
import time
import traceback
import uuid
import weakref
import zlib

from . import msgpack

PY_35 = sys.version_info >= (3, 5)
PY_352 = sys.version_info >= (3, 5, 2)

log = logging.getLogger(__name__)


@msgpack.register(object, 0x01)
class CustomEncoder:

    """Encode custom objects registered before."""

    @classmethod
    def __msgpack_encode__(cls, data, data_type):
        data_type = type(data)
        encoder = msgpack.get_custom_encoder(data_type)
        if encoder is None:
            raise TypeError("There is no custom encoder for this type registered: {}"
                            .format(data_type))
        wrapped = {
            'type': data_type.__name__,
            'module': data_type.__module__,
            'data': encoder.__msgpack_encode__(data, data_type)
        }
        return msgpack.encode(wrapped)

    @classmethod
    def __msgpack_decode__(cls, encoded_data, data_type):
        wrapped = msgpack.decode(encoded_data)
        module = sys.modules[wrapped['module']]
        data_type = getattr(module, wrapped['type'])
        encoder = msgpack.get_custom_encoder(data_type)
        if encoder is None:
            raise TypeError("There is no custom encoder for this type registered: {}"
                            .format(data_type))
        data = encoder.__msgpack_decode__(wrapped['data'], data_type)
        return data


@msgpack.register(tuple, 0x02)
class TupleEncoder:

    """Encoder for :py:obj:`tuple`."""

    @classmethod
    def __msgpack_encode__(cls, data, data_type):
        return msgpack.encode(list(data))

    @classmethod
    def __msgpack_decode__(cls, encoded_data, data_type):
        return tuple(msgpack.decode(encoded_data))


@msgpack.register(set, 0x03)
class SetEncoder:

    """Encoder for :py:obj:`set`."""

    @classmethod
    def __msgpack_encode__(cls, data, data_type):
        return msgpack.encode(list(data))

    @classmethod
    def __msgpack_decode__(cls, encoded_data, data_type):
        return set(msgpack.decode(encoded_data))


@msgpack.register(Exception, 0x04)
class ExceptionEncoder:

    """Encoder for :py:obj:`Exception`."""

    @classmethod
    def __msgpack_encode__(cls, data, data_type):
        return msgpack.encode(data.args)

    @classmethod
    def __msgpack_decode__(cls, encoded_data, data_type):
        return data_type(*msgpack.decode(encoded_data))


@msgpack.register(StopAsyncIteration, 0x05)
class StopAsyncIterationEncoder:

    """Encoder for :py:obj:`StopAsyncIteration`."""

    @classmethod
    def __msgpack_encode__(cls, data, data_type):
        return msgpack.encode(data.args)

    @classmethod
    def __msgpack_decode__(cls, encoded_data, data_type):
        return StopAsyncIteration(*msgpack.decode(encoded_data))


class Uid(uuid.UUID):

    """A unique id."""

    def __init__(self, bytes=None):     # pylint: disable=W0622
        if bytes is None:
            super().__init__(bytes=uuid.uuid1().bytes, version=1)
        else:
            super().__init__(bytes=bytes, version=1)

    @property
    def time(self):
        uid_time = (super().time - 0x01b21dd213814000) * 100 / 1e9
        return uid_time

    @classmethod
    def __msgpack_encode__(cls, data, data_type):
        return data.bytes

    @classmethod
    def __msgpack_decode__(cls, encoded, data_type):
        return cls(bytes=encoded)


class ConnectionLostStreamReaderProtocol(asyncio.StreamReaderProtocol):

    """Call a callback on connection_lost."""

    def __init__(self, *args, connection_lost_cb, **kwargs):
        super().__init__(*args, **kwargs)
        self.connection_lost_cb = connection_lost_cb

    def connection_lost(self, exc):
        super().connection_lost(exc)
        self.connection_lost_cb(exc)


class Incomming(asyncio.StreamReader):

    """A context for an incomming pipe."""

    def __init__(self, *, connection_lost_cb=None, pipe=sys.stdin, loop=None):
        super(Incomming, self).__init__(loop=loop)
        self.pipe = os.fdopen(pipe) if isinstance(pipe, int) else pipe
        self.connection_lost_cb = connection_lost_cb

    async def __aenter__(self):
        await self.connect()
        return self

    async def connect(self):
        """Connect the pipe."""
        if self.connection_lost_cb:
            protocol = ConnectionLostStreamReaderProtocol(
                self, connection_lost_cb=self.connection_lost_cb, loop=self._loop
            )
        else:
            protocol = asyncio.StreamReaderProtocol(self, loop=self._loop)

        transport, protocol = await self._loop.connect_read_pipe(
            lambda: protocol,
            self.pipe,
        )

        return transport, protocol

    async def __aexit__(self, exc_type, value, tb):
        self._transport.close()

    async def readexactly(self, n):
        """Read exactly n bytes from the stream.

        This is a short and faster implementation then the original one
        (see of https://github.com/python/asyncio/issues/394).

        """
        buffer, missing = bytearray(), n

        while missing:
            if not self._buffer:
                await self._wait_for_data('readexactly')

            if self._eof or not self._buffer:
                raise asyncio.IncompleteReadError(bytes(buffer), n)

            length = min(len(self._buffer), missing)
            buffer.extend(self._buffer[:length])

            del self._buffer[:length]
            missing -= length

            self._maybe_resume_transport()

        return buffer


class Outgoing:

    """A context for an outgoing pipe."""

    def __init__(self, *, pipe=sys.stdout, reader=None, loop=None):
        self.loop = loop if loop is not None else asyncio.get_event_loop()
        self.pipe = os.fdopen(pipe) if isinstance(pipe, int) else pipe
        self.transport = None
        self.reader = reader
        self.writer = None

    async def __aenter__(self):
        writer = await self.connect()
        return writer

    async def connect(self):
        """Connect the pipe."""
        self.transport, protocol = await self.loop.connect_write_pipe(
            asyncio.streams.FlowControlMixin,
            self.pipe
        )

        writer = asyncio.streams.StreamWriter(
            self.transport,
            protocol,
            self.reader,
            self.loop
        )
        return writer

    async def __aexit__(self, exc_type, value, tb):
        self.transport.close()


def split_data(data, size=1024):
    """Create a generator to split data into chunks."""
    data_view, data_len, start = memoryview(data), len(data), 0
    while start < data_len:
        end = min(start + size, data_len)
        yield data_view[start:end]
        start = end


class OrderedMeta(type):

    """Preserve the order of instance creation."""

    items = []

    def __call__(cls, *args, **kwargs):
        inst = super().__call__(*args, **kwargs)
        cls.items.append(inst)
        return inst

    @classmethod
    def ordered_items(mcs, dct, cls_order=None):
        """Sort and filter items by type and instance creation."""
        def _sort(item):
            if cls_order is not None:
                try:
                    cls_index = cls_order.index(item[1].__class__)
                except ValueError:
                    cls_index = None
                return cls_index, mcs.items.index(item[1])
            return mcs.items.index(item[1])

        items = collections.OrderedDict(sorted(
            (item for item in dct.items() if item[1] in mcs.items),
            key=_sort
        ))
        return items


class BaseHeaderItem(metaclass=OrderedMeta):

    """A base item of a header."""

    def __init__(self, *, default=None, encoder=None, decoder=None):
        self.index = None
        self.name = None
        self.default = default
        self.encoder = encoder
        self.decoder = decoder

    def set_index(self, index):
        """Set the index of the item in the header."""
        self.index = index

    def __set_name__(self, owner, name):
        self.name = name

    def encode(self, value):
        """Encode the value of that item."""
        if self.encoder:
            value = self.encoder(value, self)
        return value

    @staticmethod
    def decode(value):
        """Decode the value of that item."""
        return value

    def __get__(self, inst, cls):
        if inst is None:
            return self
        value = self.decode(inst)
        if self.decoder:
            value = self.decoder(value, self)
        setattr(inst, self.name, value)
        return value


class HeaderItem(BaseHeaderItem):

    """An item of a header."""

    def __init__(self, fmt, **kwargs):
        super().__init__(**kwargs)
        self.fmt = fmt

    @property
    def size(self):
        """The size of the item."""
        return struct.calcsize(self.fmt)

    def encode(self, value):
        value = super().encode(value)
        return struct.pack(self.fmt, value or self.default)

    def decode(self, value):
        value, = struct.unpack_from(self.fmt, value, self.index)
        return value


class Flag(BaseHeaderItem):

    """A boolean flag of a header."""

    def encode(self, value):
        value = super().encode(value)
        return 0 if not value else value << self.index

    def decode(self, value):
        value, = struct.unpack_from('!1B', value, 0)
        value = value >> self.index & 1
        return value


class HeaderMeta(type):

    """Order items and set the size of the header."""

    def __new__(mcs, name, bases, dct):
        items = dct['items'] = OrderedMeta.ordered_items(dct, (Flag, HeaderItem))
        dct['size'] = mcs.apply_items_index(items)
        cls = type.__new__(mcs, name, bases, dct)
        if sys.version_info < (3, 6):
            for item_name, item in items.items():
                item.__set_name__(cls, item_name)
        return cls

    @classmethod
    def apply_items_index(mcs, items):
        """Apply the index of each item."""
        flag_count = 0
        item_size = 0
        for item in items.values():
            if isinstance(item, Flag):
                item.set_index(flag_count)
                flag_count += 1
                item_size = (flag_count // 8) + (flag_count % 8 != 0)
            else:
                # items are ordered after type
                item.set_index(item_size)
                item_size += item.size
        return item_size

    def __call__(cls, encoded_bytes=None, **items):
        if encoded_bytes:
            inst = super().__call__(encoded_bytes)
            return inst

        flags = 0
        encoded_items = []
        for name, item in cls.items.items():
            if isinstance(item, Flag):
                flags += item.encode(items.get(name, False))
            else:
                encoded_items.append(item.encode(items.get(name, None)))
        encoded_flags = bytes([flags])
        encoded_items.insert(0, encoded_flags)

        inst = super().__call__(b''.join(encoded_items))
        return inst


class Header(bytes, metaclass=HeaderMeta):

    """The chunk header with flags and items."""

    eom = Flag()
    send_ack = Flag()
    recv_ack = Flag()
    compression = Flag()
    uid = HeaderItem('!16s', encoder=lambda uid, _: uid.bytes, decoder=lambda x, _: Uid(bytes=x))
    channel_name_len = HeaderItem('!H', default=0)
    data_len = HeaderItem('!I', default=0)

    def __repr__(self):
        return '<Header {self.uid} {eom} {self.data_len}>'\
            .format(self=self, eom=' eom' if self.eom else '')


Chunk = collections.namedtuple('Chunk', ('header', 'channel_name', 'data'))


class Channels:

    """Hold references to all channel queues."""
{'fraction_non_alphanumeric': 0.10378045125658261, 'fraction_numerical': 0.007107077107505245, 'mean_word_length': 2.4522613065326633, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 13, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '734015', 'n_tokens_mistral': 7008, 'n_tokens_neox': 6674, 'n_words': 1754}