<commit_before>from functools import partial
from datetime import timedelta

from tornado import gen
from tornadoredis import Client
from tornadoredis.exceptions import ResponseError
from tornadoredis.pubsub import BaseSubscriber


class CelerySubscriber(BaseSubscriber):
    def unsubscribe_channel(self, channel_name):
        """Unsubscribes the redis client from the channel"""
        del self.subscribers[channel_name]
        del self.subscriber_count[channel_name]
        self.redis.unsubscribe(channel_name)

    def on_message(self, msg):
        if not msg:
            return
        try:
            if msg.kind == 'message' and msg.body:
                # Get the list of subscribers for this channel
                for subscriber in self.subscribers[msg.channel].keys():
                    subscriber(msg.body)
        finally:
            super(CelerySubscriber, self).on_message(msg)


class RedisClient(Client):
    @gen.engine
    def _consume_bulk(self, tail, callback=None):
        response = yield gen.Task(self.connection.read, int(tail) + 2)
        if isinstance(response, Exception):
            raise response
        if not response:
            raise ResponseError('EmptyResponse')
        else:
            # We don't cast try to convert to unicode here as the response
            # may not be utf-8 encoded, for example if using msgpack as a
            # serializer
            # response = to_unicode(response)
            response = response[:-2]
        callback(response)


class RedisConsumer(object):
    def __init__(self, producer):
        self.producer = producer
        backend = producer.app.backend
        self.client = RedisClient(host=backend.connparams['host'],
                                  port=backend.connparams['port'],
                                  password=backend.connparams['password'],
                                  selected_db=backend.connparams['db'],
                                  io_loop=producer.conn_pool.io_loop)
        self.client.connect()
        self.subscriber = CelerySubscriber(self.client)

    def wait_for(self, task_id, callback, expires=None, persistent=None):
        key = self.producer.app.backend.get_key_for_task(task_id)
        if expires:
            timeout = self.producer.conn_pool.io_loop.add_timeout(
                timedelta(milliseconds=expires), self.on_timeout, key)
        else:
            timeout = None
        self.subscriber.subscribe(
            key, partial(self.on_result, key, callback, timeout))

    def on_result(self, key, callback, timeout, result):
        if timeout:
            self.producer.conn_pool.io_loop.remove_timeout(timeout)
        self.subscriber.unsubscribe_channel(key)
        callback(result)

    def on_timeout(self, key):
        self.subscriber.unsubscribe_channel(key)<commit_msg>call result callback after redis has unsubscribed<commit_after>from functools import partial
from datetime import timedelta

from tornado import gen
from tornadoredis import Client
from tornadoredis.exceptions import ResponseError
from tornadoredis.pubsub import BaseSubscriber


class CelerySubscriber(BaseSubscriber):
    def unsubscribe_channel(self, channel_name, callback=None):
        """Unsubscribes the redis client from the channel"""
        del self.subscribers[channel_name]
        del self.subscriber_count[channel_name]
        self.redis.unsubscribe(channel_name, callback)

    def on_message(self, msg):
        if not msg:
            return
        try:
            if msg.kind == 'message' and msg.body:
                # Get the list of subscribers for this channel
                for subscriber in self.subscribers[msg.channel].keys():
                    subscriber(msg.body)
        finally:
            super(CelerySubscriber, self).on_message(msg)


class RedisClient(Client):
    @gen.engine
    def _consume_bulk(self, tail, callback=None):
        response = yield gen.Task(self.connection.read, int(tail) + 2)
        if isinstance(response, Exception):
            raise response
        if not response:
            raise ResponseError('EmptyResponse')
        else:
            # We don't cast try to convert to unicode here as the response
            # may not be utf-8 encoded, for example if using msgpack as a
            # serializer
            # response = to_unicode(response)
            response = response[:-2]
        callback(response)


class RedisConsumer(object):
    def __init__(self, producer):
        self.producer = producer
        backend = producer.app.backend
        self.client = RedisClient(host=backend.connparams['host'],
                                  port=backend.connparams['port'],
                                  password=<PASSWORD>.connparams['password'],
                                  selected_db=backend.connparams['db'],
                                  io_loop=producer.conn_pool.io_loop)
        self.client.connect()
        self.subscriber = CelerySubscriber(self.client)

    def wait_for(self, task_id, callback, expires=None, persistent=None):
        key = self.producer.app.backend.get_key_for_task(task_id)
        if expires:
            timeout = self.producer.conn_pool.io_loop.add_timeout(
                timedelta(milliseconds=expires), self.on_timeout, key)
        else:
            timeout = None
        self.subscriber.subscribe(
            key, partial(self.on_result, key, callback, timeout))

    def on_result(self, key, callback, timeout, result):
        if timeout:
            self.producer.conn_pool.io_loop.remove_timeout(timeout)
        self.subscriber.unsubscribe_channel(key, lambda: callback(result))

    def on_timeout(self, key):
        self.subscriber.unsubscribe_channel(key)
{'fraction_non_alphanumeric': 0.08302872062663186, 'fraction_numerical': 0.0010443864229765013, 'mean_word_length': 2.3213872832369944, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5153614', 'n_tokens_mistral': 1560, 'n_tokens_neox': 1508, 'n_words': 368}