<commit_before>from kombu.connection import Resource
from kombu.messaging import Producer

from itertools import chain

__all__ = ["ProducerPool", "connections", "producers", "set_limit", "reset"]
_limit = [200]
_groups = []


def register_group(group):
    _groups.append(group)
    return group


class ProducerPool(Resource):
    Producer = Producer

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

    def create_producer(self):
        conn = self.connections.acquire(block=True)
        producer = self.Producer(conn)
        producer.connection = conn
        return producer

    def new(self):
        return lambda: self.create_producer()

    def setup(self):
        if self.limit:
            for _ in xrange(self.limit):
                self._resource.put_nowait(self.new())

    def prepare(self, p):
        if callable(p):
            p = p()
        if not p.connection:
            p.connection = self.connections.acquire(block=True)
            p.revive(p.connection.default_channel)
        return p

    def release(self, resource):
        resource.connection.release()
        resource.connection = None
        super(ProducerPool, self).release(resource)


class HashingDict(dict):

    def __getitem__(self, key):
        h = hash(key)
        if h not in self:
            return self.__missing__(key)
        return dict.__getitem__(self, h)

    def __setitem__(self, key, value):
        return dict.__setitem__(self, hash(key), value)

    def __delitem__(self, key):
        return dict.__delitem__(self, hash(key))


class PoolGroup(HashingDict):

    def create(self, resource, limit):
        raise NotImplementedError("PoolGroups must define ``create``")

    def __missing__(self, resource):
        k = self[resource] = self.create(resource, get_limit())
        return k


class _Connections(PoolGroup):

    def create(self, connection, limit):
        return connection.Pool(limit=limit)
connections = register_group(_Connections())


class _Producers(HashingDict):

    def create(self, connection, limit):
        return ProducerPool(connections[connection], limit=limit)
producers = register_group(_Producers())


def _all_pools():
    return chain(*[(g.itervalues() if g else iter([])) for g in _groups])


def get_limit():
    return _limit[0]


def set_limit(limit):
    if _limit[0] != limit:
        _limit[0] = limit
        for pool in _all_pools():
            pool.limit = limit
        reset()
    return limit


def reset(*args, **kwargs):
    for pool in _all_pools():
        try:
            pool.force_close_all()
        except Exception:
            pass
    for group in _groups:
        group.clear()


try:
    from multiprocessing.util import register_after_fork
    register_after_fork(connections, reset)
except ImportError:
    pass
<commit_msg>ProducerPool.Producer: can be used to override the producer type used<commit_after>from kombu.connection import Resource
from kombu.messaging import Producer

from itertools import chain

__all__ = ["ProducerPool", "connections", "producers", "set_limit", "reset"]
_limit = [200]
_groups = []


def register_group(group):
    _groups.append(group)
    return group


class ProducerPool(Resource):

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

    def Producer(self, connection):
        return Producer(connection)

    def create_producer(self):
        return self.Producer(self.connections.acquire(block=True))

    def new(self):
        return lambda: self.create_producer()

    def setup(self):
        if self.limit:
            for _ in xrange(self.limit):
                self._resource.put_nowait(self.new())

    def prepare(self, p):
        if callable(p):
            p = p()
        if not p.connection:
            p.connection = self.connections.acquire(block=True)
            p.revive(p.connection.default_channel)
        return p

    def release(self, resource):
        resource.connection.release()
        resource.connection = None
        super(ProducerPool, self).release(resource)


class HashingDict(dict):

    def __getitem__(self, key):
        h = hash(key)
        if h not in self:
            return self.__missing__(key)
        return dict.__getitem__(self, h)

    def __setitem__(self, key, value):
        return dict.__setitem__(self, hash(key), value)

    def __delitem__(self, key):
        return dict.__delitem__(self, hash(key))


class PoolGroup(HashingDict):

    def create(self, resource, limit):
        raise NotImplementedError("PoolGroups must define ``create``")

    def __missing__(self, resource):
        k = self[resource] = self.create(resource, get_limit())
        return k


class _Connections(PoolGroup):

    def create(self, connection, limit):
        return connection.Pool(limit=limit)
connections = register_group(_Connections())


class _Producers(HashingDict):

    def create(self, connection, limit):
        return ProducerPool(connections[connection], limit=limit)
producers = register_group(_Producers())


def _all_pools():
    return chain(*[(g.itervalues() if g else iter([])) for g in _groups])


def get_limit():
    return _limit[0]


def set_limit(limit):
    if _limit[0] != limit:
        _limit[0] = limit
        for pool in _all_pools():
            pool.limit = limit
        reset()
    return limit


def reset(*args, **kwargs):
    for pool in _all_pools():
        try:
            pool.force_close_all()
        except Exception:
            pass
    for group in _groups:
        group.clear()


try:
    from multiprocessing.util import register_after_fork
    register_after_fork(connections, reset)
except ImportError:
    pass

{'fraction_non_alphanumeric': 0.10563139931740614, 'fraction_numerical': 0.0020477815699658703, 'mean_word_length': 2.8660949868073877, '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': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6460207', 'n_tokens_mistral': 1764, 'n_tokens_neox': 1684, 'n_words': 440}