handler.py 1.3 KB
import json
import time
import datetime


class DateTimeJSONEncoder(json.JSONEncoder):
    """
    JSONEncoder subclass that knows how to encode date/time, decimal types, and
    UUIDs.
    """
    def default(self, o):
        # See "Date Time String Format" in the ECMA-262 specification.
        if isinstance(o, datetime.datetime):
            r = o.isoformat()
            if o.microsecond:
                r = r[:23] + r[26:]
            if r.endswith('+00:00'):
                r = r[:-6] + 'Z'
            return r
        else:
            return super().default(o)


def dict_str_to_int(res):
    for k, v in res.items():
        res[k] = int(v)


class RedisHandler:

    def __init__(self, redis):
        self.redis = redis
        self.time_expires = datetime.timedelta(hours=24)
        self.time_format = '%a %b %d %H:%M:%S %Y'
        self.prefix = 'bwm_ocr'
        self.queue_key = '{0}:queue'.format(self.prefix)

    def enqueue(self, task_id):
        # 1
        mapping = {task_id: time.time()}
        return self.redis.zadd(self.queue_key, mapping)

    def dequeue(self):
        # model_id:int or None
        res_list = self.redis.zremrangebyrank(self.queue_key, 0, 0)
        pop_item_list = res_list[0]
        pop_item = int(pop_item_list[0]) if pop_item_list else None
        return pop_item