handler.py
1.3 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
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