<commit_before>import asyncio
import json
import logging
import pickle
import time
from enum import Enum
from functools import partial
from typing import Any, Callable, Dict, Hashable, Optional, TypeVar

from aiormq.tools import awaitable

from aio_pika.channel import Channel
from aio_pika.exceptions import DeliveryError
from aio_pika.exchange import ExchangeType
from aio_pika.message import (
    DeliveryMode, IncomingMessage, Message, ReturnedMessage,
)
from aio_pika.tools import shield

from .base import Base, Proxy


log = logging.getLogger(__name__)

R = TypeVar("R")
P = TypeVar("P")
CallbackType = Callable[[P], R]


class RPCMessageTypes(Enum):
    error = "error"
    result = "result"
    call = "call"


class RPC(Base):
    __slots__ = (
        "channel",
        "loop",
        "proxy",
        "result_queue",
        "result_consumer_tag",
        "routes",
        "consumer_tags",
        "dlx_exchange",
    )

    DLX_NAME = "rpc.dlx"
    DELIVERY_MODE = DeliveryMode.NOT_PERSISTENT

    __doc__ = """
    Remote Procedure Call helper.

    Create an instance ::

        rpc = await RPC.create(channel)

    Registering python function ::

        # RPC instance passes only keyword arguments
        def multiply(*, x, y):
            return x * y

        await rpc.register("multiply", multiply)

    Call function through proxy ::

        assert await rpc.proxy.multiply(x=2, y=3) == 6

    Call function explicit ::

        assert await rpc.call('multiply', dict(x=2, y=3)) == 6

    """

    def __init__(self, channel: Channel):
        self.channel = channel
        self.loop = self.channel.loop
        self.proxy = Proxy(self.call)
        self.result_queue = None
        self.futures = {}  # type Dict[int, asyncio.Future]
        self.result_consumer_tag = None
        self.routes = {}
        self.queues = {}
        self.consumer_tags = {}
        self.dlx_exchange = None

    def __remove_future(self, future: asyncio.Future):
        log.debug("Remove done future %r", future)
        self.futures.pop(id(future), None)

    def create_future(self) -> asyncio.Future:
        future = self.loop.create_future()
        log.debug("Create future for RPC call")

        self.futures[id(future)] = future
        future.add_done_callback(self.__remove_future)
        return future

    @shield
    async def close(self):
        if self.result_queue is None:
            log.warning("RPC already closed")
            return

        log.debug("Cancelling listening %r", self.result_queue)
        await self.result_queue.cancel(self.result_consumer_tag)
        self.result_consumer_tag = None

        log.debug("Unbinding %r", self.result_queue)
        await self.result_queue.unbind(
            self.dlx_exchange,
            "",
            arguments={"From": self.result_queue.name, "x-match": "any"},
        )

        log.debug("Cancelling undone futures %r", self.futures)
        for future in self.futures.values():
            if future.done():
                continue

            future.set_exception(asyncio.CancelledError)

        log.debug("Deleting %r", self.result_queue)
        await self.result_queue.delete()
        self.result_queue = None

    @shield
    async def initialize(self, auto_delete=True, durable=False, **kwargs):
        if self.result_queue is not None:
            return

        self.result_queue = await self.channel.declare_queue(
            None, auto_delete=auto_delete, durable=durable, **kwargs
        )

        self.dlx_exchange = await self.channel.declare_exchange(
            self.DLX_NAME, type=ExchangeType.HEADERS, auto_delete=True,
        )

        await self.result_queue.bind(
            self.dlx_exchange,
            "",
            arguments={"From": self.result_queue.name, "x-match": "any"},
        )

        self.result_consumer_tag = await self.result_queue.consume(
            self.on_result_message, exclusive=True, no_ack=True,
        )

        self.channel.add_close_callback(self.on_close)
        self.channel.add_on_return_callback(
            self.on_message_returned, weak=False,
        )

    def on_close(self, exc=None):
        log.debug("Closing RPC futures because %r", exc)
        for future in self.futures.values():<commit_msg>Fixed callback on_close for rpc<commit_after>import asyncio
import json
import logging
import pickle
import time
from enum import Enum
from functools import partial
from typing import Any, Callable, Dict, Hashable, Optional, TypeVar

from aiormq.tools import awaitable

from aio_pika.channel import Channel
from aio_pika.exceptions import DeliveryError
from aio_pika.exchange import ExchangeType
from aio_pika.message import (
    DeliveryMode, IncomingMessage, Message, ReturnedMessage,
)
from aio_pika.tools import shield

from .base import Base, Proxy


log = logging.getLogger(__name__)

R = TypeVar("R")
P = TypeVar("P")
CallbackType = Callable[[P], R]


class RPCMessageTypes(Enum):
    error = "error"
    result = "result"
    call = "call"


class RPC(Base):
    __slots__ = (
        "channel",
        "loop",
        "proxy",
        "result_queue",
        "result_consumer_tag",
        "routes",
        "consumer_tags",
        "dlx_exchange",
    )

    DLX_NAME = "rpc.dlx"
    DELIVERY_MODE = DeliveryMode.NOT_PERSISTENT

    __doc__ = """
    Remote Procedure Call helper.

    Create an instance ::

        rpc = await RPC.create(channel)

    Registering python function ::

        # RPC instance passes only keyword arguments
        def multiply(*, x, y):
            return x * y

        await rpc.register("multiply", multiply)

    Call function through proxy ::

        assert await rpc.proxy.multiply(x=2, y=3) == 6

    Call function explicit ::

        assert await rpc.call('multiply', dict(x=2, y=3)) == 6

    """

    def __init__(self, channel: Channel):
        self.channel = channel
        self.loop = self.channel.loop
        self.proxy = Proxy(self.call)
        self.result_queue = None
        self.futures = {}  # type Dict[int, asyncio.Future]
        self.result_consumer_tag = None
        self.routes = {}
        self.queues = {}
        self.consumer_tags = {}
        self.dlx_exchange = None

    def __remove_future(self, future: asyncio.Future):
        log.debug("Remove done future %r", future)
        self.futures.pop(id(future), None)

    def create_future(self) -> asyncio.Future:
        future = self.loop.create_future()
        log.debug("Create future for RPC call")

        self.futures[id(future)] = future
        future.add_done_callback(self.__remove_future)
        return future

    @shield
    async def close(self):
        if self.result_queue is None:
            log.warning("RPC already closed")
            return

        log.debug("Cancelling listening %r", self.result_queue)
        await self.result_queue.cancel(self.result_consumer_tag)
        self.result_consumer_tag = None

        log.debug("Unbinding %r", self.result_queue)
        await self.result_queue.unbind(
            self.dlx_exchange,
            "",
            arguments={"From": self.result_queue.name, "x-match": "any"},
        )

        log.debug("Cancelling undone futures %r", self.futures)
        for future in self.futures.values():
            if future.done():
                continue

            future.set_exception(asyncio.CancelledError)

        log.debug("Deleting %r", self.result_queue)
        await self.result_queue.delete()
        self.result_queue = None

    @shield
    async def initialize(self, auto_delete=True, durable=False, **kwargs):
        if self.result_queue is not None:
            return

        self.result_queue = await self.channel.declare_queue(
            None, auto_delete=auto_delete, durable=durable, **kwargs
        )

        self.dlx_exchange = await self.channel.declare_exchange(
            self.DLX_NAME, type=ExchangeType.HEADERS, auto_delete=True,
        )

        await self.result_queue.bind(
            self.dlx_exchange,
            "",
            arguments={"From": self.result_queue.name, "x-match": "any"},
        )

        self.result_consumer_tag = await self.result_queue.consume(
            self.on_result_message, exclusive=True, no_ack=True,
        )

        self.channel.add_close_callback(self.on_close)
        self.channel.add_on_return_callback(
            self.on_message_returned, weak=False,
        )

    def on_close(self, _: Channel, exc: Optional[BaseException] = None):
        log.debug("Closing RPC futures because %r", exc)
        for future in self.futures.values():
{'fraction_non_alphanumeric': 0.10154346060113728, 'fraction_numerical': 0.0013925960311013113, 'mean_word_length': 2.4499599679743795, 'pattern_counts': {'":': 8, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, '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': '5899133', 'n_tokens_mistral': 2666, 'n_tokens_neox': 2557, 'n_words': 594}