class CompositeClosableAsyncObject(ClosableAsyncObject):
    """
    Base class for objects that can be asynchronously closed and awaited and
    have ownership of other closable objects.
    """
    def on_open(self):
        self._children = set()

    async def on_close(self):
        await asyncio.gather(
            *[
                self.on_close_child(child)
                for child in self._children
            ],
            loop=self.loop,
        )
        del self._children

    async def on_close_child(self, child):
        """
        Called whenever a child instance must be closed.

        :returns: An awaitable that must only complete when the specified child
            instance is effectively closed.

        May be redefined by child-classes. The default implementation calls
            `close` followed by `wait_closed` on the specified child instance.
        """
        child.close()
        await child.wait_closed()

    def register_child(self, child):
        """
        Register a new child that will be closed whenever the current instance
        closes.

        :param child: The child instance.
        """
        if self.closing:
            child.close()
        else:
            self._children.add(child)
            child.on_closed.connect(self.unregister_child)

    def unregister_child(self, child):
        """
        Unregister an existing child that is no longer to be owned by the
        current instance.

        :param child: The child instance.
        """
        self._children.remove(child)
        child.on_closed.disconnect(self.unregister_child) def cancel_on_closing(func):
    """
    Automatically cancels a function or coroutine when the defining instance
    gets closed.

    :param func: The function to cancel on closing.
    :returns: A decorated function or coroutine.
    """
    @wraps(func)
    async def wrapper(self, *args, **kwargs):
        return await self.await_until_closing(func(self, *args, **kwargs))

    return wrapper class AsyncBox(ClosableAsyncObject):
    def on_open(self, maxsize=0):
        self._maxsize = maxsize
        self._queue = asyncio.Queue(maxsize=maxsize, loop=self.loop)
        self._can_read = asyncio.Event(loop=self.loop)
        self._can_write = asyncio.Event(loop=self.loop)
        self._can_write.set()

    @cancel_on_closing
    async def read(self):
        """
        Read from the box in a blocking manner.

        :returns: An item from the box.
        """
        result = await self._queue.get()

        self._can_write.set()

        if self._queue.empty():
            self._can_read.clear()

        return result

    def read_nowait(self):
        """
        Read from the box in a non-blocking manner.

        If the box is empty, an exception is thrown. You should always check
        for emptiness with `empty` or `wait_not_empty` before calling this
        method.

        :returns: An item from the box.
        """
        result = self._queue.get_nowait()

        self._can_write.set()

        if self._queue.empty():
            self._can_read.clear()

        return result

    @cancel_on_closing
    async def wait_not_empty(self):
        """
        Blocks until the queue is not empty.
        """
        await self._can_read.wait()

    @cancel_on_closing
    async def wait_not_full(self):
        """
        Blocks until the queue is not full.
        """
        await self._can_write.wait()

    def full(self):
        return self._queue.full()

    def empty(self):
        return self._queue.empty()

    @cancel_on_closing
    async def write(self, item):
        """
        Write an item in the queue.

        :param item: The item.
        """
        await self._queue.put(item)
        self._can_read.set()

        if self._queue.full():
            self._can_write.clear()

    def write_nowait(self, item):
        """
        Write in the box in a non-blocking manner.

        If the box is full, an exception is thrown. You should always check
        for fullness with `full` or `wait_not_full` before calling this method.

        :param item: An item.
        """
        self._queue.put_nowait(item)
        self._can_read.set()

        if self._queue.full():
            self._can_write.clear()

    def clear(self):
        """
        Clear the box.
        """
        while not self._queue.empty():
            self._queue.get_nowait()

        self._can_read.clear()
        self._can_write.set()

    def clone(self):
        """
        Clone the box.

        :returns: A new box with the same item queue.

        The cloned box is not closed, no matter the initial state of the
        original instance.
        """
        result = AsyncBox(maxsize=self._maxsize, loop=self.loop)
        result._queue = self._queue
        result._can_read = self._can_read
        result._can_write = self._can_write

        return result DEALER = b'DEALER' PAIR = b'PAIR' PUB = b'PUB' PUSH = b'PUSH' PULL = b'PULL' REP = b'REP' REQ = b'REQ' ROUTER = b'ROUTER' SUB = b'SUB' XPUB = b'XPUB' XSUB = b'XSUB' class AsyncList(list, AsyncObject):
    """
    A list-like object that is asyncio-compatible.
    """
    def __init__(self, *items, **kwargs):
        """
        Instanciate a new `AsyncList` list with optional items.

        :param items: If provided, an iterable of items to fill the list with.
        :param loop: The event loop the instance is associated with. Default
            to the thread-local current asyncio event-loop.
        """
        list.__init__(self, *items)
        AsyncObject.__init__(self, **kwargs)
        self._empty = asyncio.Event(loop=self.loop)
        self._not_empty = asyncio.Event(loop=self.loop)
        self._change_futures = set()
        self._refresh()

    def _refresh(self):
        if self:
            self._empty.clear()
            self._not_empty.set()
        else:
            self._empty.set()
            self._not_empty.clear()

        for future in self._change_futures:
            if not future.done():
                future.set_result(None)

    def append(self, *args, **kwargs):
        super().append(*args, **kwargs)
        self._refresh()

    def remove(self, *args, **kwargs):
        super().remove(*args, **kwargs)
        self._refresh()

    def extend(self, *args, **kwargs):
        super().extend(*args, **kwargs)
        self._refresh()

    def clear(self, *args, **kwargs):
        super().clear(*args, **kwargs)
        self._refresh()

    def __setitem__(self, *args, **kwargs):
        super().__setitem__(*args, **kwargs)
        self._refresh()

    def pop(self, *args, **kwargs):
        try:
            return super().pop(*args, **kwargs)
        finally:
            self._refresh()

    def insert(self, *args, **kwargs):
        super().insert(*args, **kwargs)
        self._refresh()

    async def wait_not_empty(self):
        """
        Wait for the list to become non-empty.
        """
        await self._not_empty.wait()

    async def wait_empty(self):
        """
        Wait for the list to become empty.
        """
        await self._empty.wait()

    async def wait_change(self):
        """
        Wait for the list to change.
        """
        future = asyncio.Future(loop=self.loop)
        self._change_futures.add(future)
        future.add_done_callback(self._change_futures.discard)
        await future

    def create_proxy(self):
        """
        Create a fair list proxy to the current instance.

        :returns: A `FairListProxy` instance.
        """
        return FairListProxy(self) class UnsupportedSchemeError(AZMQError):
    def __init__(self, scheme):
        super().__init__("Unsupported scheme '%s'" % scheme)
        self.scheme = scheme class InvalidOperation(AZMQError):
    def __init__(self, msg):
        super().__init__(msg) class TCPClientEngine(BaseEngine):
    def on_open(self, *, host, port, **kwargs):
        super().on_open(**kwargs)

        self.host = host
        self.port = port

    async def open_connection(self):
        reader, writer = await asyncio.open_connection(
            host=self.host,
            port=self.port,
            loop=self.loop,
        )

        writer.transport.get_extra_info('socket').setsockopt(
            socket.IPPROTO_TCP,
            socket.TCP_NODELAY,
            1,
        )

        address, port = writer.get_extra_info('peername')
        logger.debug(
            "Connection to %s:%s established.",
            address,
            port,
        )

        try:
            async with StreamConnection(
                reader=reader,
                writer=writer,
                address=address,
                zap_client=self.zap_client,
                socket_type=self.socket_type,
                identity=self.identity,
                mechanism=self.mechanism,
                on_ready=self.on_connection_ready.emit,
                on_lost=self.on_connection_lost.emit,
                on_failure=self.on_connection_failure,
                loop=self.loop,
            ) as connection:
                self.register_child(connection)
                await connection.wait_closed()

            return connection.get_error()
        finally:
            logger.debug(
                "Connection to %s:%s closed.",
                address,
                port,
            ) class TCPServerEngine(BaseEngine):
    def on_open(self, *, host, port, **kwargs):
        super().on_open(**kwargs)

        self.host = host
        self.port = port

    async def open_connection(self):
        try:
            server = await asyncio.start_server(
                self.handle_connection,
                host=self.host,
                port=self.port,
                loop=self.loop,
            )

            try:
                await server.wait_closed()
            finally:
                server.close()
                await server.wait_closed()

        except asyncio.CancelledError:
            raise
        except Exception:
            logger.exception(
                "Unable to start TCP server on %s:%s.",
                self.host,
                self.port,
            )

    async def handle_connection(self, reader, writer):
        address, port = writer.get_extra_info('peername')
        peername = '%s:%s' % (address, port)

        logger.debug("Connection from %s established.", peername)
        writer.transport.get_extra_info('socket').setsockopt(
            socket.IPPROTO_TCP,
            socket.TCP_NODELAY,
            1,
        )

        async with StreamConnection(
            reader=reader,
            writer=writer,
            address=address,
            zap_client=self.zap_client,
            socket_type=self.socket_type,
            identity=self.identity,
            mechanism=self.mechanism,
            on_ready=self.on_connection_ready.emit,
            on_lost=self.on_connection_lost.emit,
            on_failure=self.on_connection_failure,
            loop=self.loop,
        ) as connection:
            self.register_child(connection)
            await connection.wait_closed()

        logger.debug("Connection from %s lost.", peername) class IPCClientEngine(BaseEngine):
    def on_open(self, *, path, **kwargs):
        super().on_open(**kwargs)

        self.path = path

    async def open_connection(self):
        reader, writer = await open_ipc_connection(
            path=self.path,
            loop=self.loop,
        )

        logger.debug("Connection to %s established.", self.path)

        try:
            async with StreamConnection(
                reader=reader,
                writer=writer,
                address=self.path,
                zap_client=self.zap_client,
                socket_type=self.socket_type,
                identity=self.identity,
                mechanism=self.mechanism,
                on_ready=self.on_connection_ready.emit,
                on_lost=self.on_connection_lost.emit,
                on_failure=self.on_connection_failure,
                loop=self.loop,
            ) as connection:
                self.register_child(connection)
                await connection.wait_closed()

            return connection.get_error()
        finally:
            logger.debug("Connection to %s closed.", self.path) class IPCServerEngine(BaseEngine):
    def on_open(self, *, path, **kwargs):
        super().on_open(**kwargs)

        self.path = path

    async def open_connection(self):
        try:
            server = await start_ipc_server(
                self.handle_connection,
                path=self.path,
                loop=self.loop,
            )

            try:
                await server.wait_closed()
            finally:
                server.close()
                await server.wait_closed()

        except asyncio.CancelledError:
            raise
        except Exception:
            logger.exception(
                "Unable to start UNIX server on %s.",
                self.path,
            )

    async def handle_connection(self, reader, writer):
        logger.debug("Connection from %s established.", self.path)

        async with StreamConnection(
            reader=reader,
            writer=writer,
            address=self.path,
            zap_client=self.zap_client,
            socket_type=self.socket_type,
            identity=self.identity,
            mechanism=self.mechanism,
            on_ready=self.on_connection_ready.emit,
            on_lost=self.on_connection_lost.emit,
            on_failure=self.on_connection_failure,
            loop=self.loop,
        ) as connection:
            self.register_child(connection)
            await connection.wait_closed()

        logger.debug("Connection from %s lost.", self.path) class InprocClientEngine(BaseEngine):
    def on_open(self, *, context, path, **kwargs):
        super().on_open(**kwargs)

        self.context = context
        self.path = path

    async def open_connection(self):
        try:
            channel = await self.await_until_closing(
                self.context._open_inproc_connection(
                    path=self.path,
                ),
            )
        except Exception as ex:
            logger.debug("Connection to %s failed: %s", self.path, ex)
        else:
            logger.debug(
                "Connection to %s established.",
                self.path,
            )

            try:
                async with channel:
                    async with InprocConnection(
                        channel=channel,
                        socket_type=self.socket_type,
                        identity=self.identity,
                        mechanism=self.mechanism,
                        on_ready=self.on_connection_ready.emit,
                        on_lost=self.on_connection_lost.emit,
                        on_failure=self.on_connection_failure,
                    ) as connection:
                        self.register_child(connection)
                        await connection.wait_closed()

                    return connection.get_error()
            finally:
                logger.debug(
                    "Connection to %s closed.",
                    self.path,
                ) class InprocServerEngine(BaseEngine):
    def on_open(self, context, path):
        super().on_open()

        self.context = context
        self.path = path

    async def open_connection(self):
        try:
            server = await self.context._start_inproc_server(
                self.handle_connection,
                path=self.path,
            )

            try:
                await server.wait_closed()
            finally:
                server.close()
                await server.wait_closed()

        except asyncio.CancelledError:
            raise
        except Exception:
            logger.exception(
                "Unable to start INPROC server on %s.",
                self.path,
            )

    async def handle_connection(self, channel):
        logger.debug("Connection from %s established.", channel.path)

        async with channel:
            async with InprocConnection(
                channel=channel,
                socket_type=self.socket_type,
                identity=self.identity,
                mechanism=self.mechanism,
                on_ready=self.on_connection_ready.emit,
                on_lost=self.on_connection_lost.emit,
                on_failure=self.on_connection_failure,
            ) as connection:
                self.register_child(connection)
                await connection.wait_closed()

        logger.debug("Connection from %s lost.", channel.path) class Null(object):
    def __call__(self):
        return NullMechanism()
import asyncio
import random
import struct
from pyslot import Signal
from urllib.parse import (
    urlsplit,
    urlunsplit,
)
from contextlib import ExitStack
from functools import partial
from .common import (
    CompositeClosableAsyncObject,
    cancel_on_closing,
    AsyncBox,
)
from .constants import (
    DEALER,
    PAIR,
    PUB,
    PUSH,
    PULL,
    REP,
    REQ,
    ROUTER,
    SUB,
    XPUB,
    XSUB,
)
from .containers import AsyncList
from .errors import (
    UnsupportedSchemeError,
    InvalidOperation,
)
from .engines.tcp import (
    TCPClientEngine,
    TCPServerEngine,
)
from .engines.ipc import (
    IPCClientEngine,
    IPCServerEngine,
)
from .engines.inproc import (
    InprocClientEngine,
    InprocServerEngine,
)
from .log import logger
from .mechanisms import Null
"""
A ZMQ socket class implementation.
"""






class Peer(object):
    """
    Represents a peer connection, either establish or in the process of being
    established.
    """
    __slots__ = [
        'engine',
        'connection',
        'inbox',
        'outbox',
    ]

    def __init__(self, engine, connection, inbox, outbox):
        self.engine = engine
        self.connection = connection
        self.inbox = inbox
        self.outbox = outbox

    def reset(self):
        self.connection = None
        self.inbox.close()
        self.inbox = self.inbox.clone()
        self.outbox.close()
        self.outbox = self.outbox.clone()

    @property
    def dead(self):
        return not self.engine and self.inbox.empty()


class Socket(CompositeClosableAsyncObject):
    """
    A ZMQ socket.

    This class is **NOT** thread-safe.
    """
    def on_open(self, context, socket_type, identity=None, mechanism=None):
        super().on_open()

        self.context = context
        self.socket_type = socket_type
        self.identity = identity
        self.mechanism = mechanism or Null()
        self.max_inbox_size = 0
        self.max_outbox_size = 0
        self._outgoing_engines = {}
        self._outgoing_peers = {}
        self._incoming_engines = {}
        self._peers = AsyncList(loop=self.loop)
        self._in_peers = self._peers.create_proxy()
        self._out_peers = self._peers.create_proxy()
        self._base_identity = random.getrandbits(32)
        self._subscriptions = []
        self._read_lock = asyncio.Lock(loop=self.loop)

        if self.socket_type == REQ:
            # This future holds the last connection we sent a request to (or
            # none, if no request was sent yet). This allows to start receiving
            # before we send.
            self._current_peer = asyncio.Future(loop=self.loop)
            self.recv_multipart = self._recv_req
            self.send_multipart = self._send_req
        elif self.socket_type == REP:
            # This future holds the last connection we received a request from
            # (or none, if no request was received yet). This allows to start
            # receiving before we send.
            self._current_peer = asyncio.Future(loop=self.loop)
            self.recv_multipart = self._recv_rep
            self.send_multipart = self._send_rep
        elif self.socket_type == DEALER:
            self.recv_multipart = self._recv_dealer
            self.send_multipart = self._send_dealer
        elif self.socket_type == ROUTER:
            self.recv_multipart = self._recv_router
            self.send_multipart = self._send_router
        elif self.socket_type == PUB:
            self.recv_multipart = self._no_recv
            self.send_multipart = self._send_pub
        elif self.socket_type == XPUB:
            self.recv_multipart = self._recv_xpub
            self.send_multipart = self._send_pub  # This is not a typo.
        elif self.socket_type == SUB:
            self.recv_multipart = self._recv_sub
            self.send_multipart = self._no_send
        elif self.socket_type == XSUB:
            self.recv_multipart = self._recv_sub  # This is not a typo.
            self.send_multipart = self._send_xsub
        elif self.socket_type == PUSH:
            self.recv_multipart = self._no_recv
            self.send_multipart = self._send_push
        elif self.socket_type == PULL:
            self.recv_multipart = self._recv_pull
            self.send_multipart = self._no_send

        elif self.socket_type == PAIR: