LibreTranslate/libretranslate/storage.py

167 lines
4.2 KiB
Python
Raw Normal View History

2023-03-10 03:00:27 +00:00
import redis
storage = None
def get_storage():
return storage
class Storage:
2023-03-11 17:06:36 +00:00
def exists(self, key):
raise Exception("not implemented")
2023-03-10 03:00:27 +00:00
def set_bool(self, key, value):
raise Exception("not implemented")
def get_bool(self, key):
raise Exception("not implemented")
def set_int(self, key, value):
raise Exception("not implemented")
def get_int(self, key):
raise Exception("not implemented")
def set_str(self, key, value):
raise Exception("not implemented")
def get_str(self, key):
raise Exception("not implemented")
2023-03-10 04:07:12 +00:00
def set_hash_int(self, ns, key, value):
2023-03-10 03:00:27 +00:00
raise Exception("not implemented")
def get_hash_int(self, ns, key):
raise Exception("not implemented")
def inc_hash_int(self, ns, key):
raise Exception("not implemented")
def dec_hash_int(self, ns, key):
raise Exception("not implemented")
2023-03-10 03:00:27 +00:00
def get_hash_keys(self, ns):
raise Exception("not implemented")
def del_hash(self, ns, key):
raise Exception("not implemented")
2023-03-10 03:00:27 +00:00
class MemoryStorage(Storage):
def __init__(self):
self.store = {}
2023-03-11 17:06:36 +00:00
def exists(self, key):
return key in self.store
2023-03-10 03:00:27 +00:00
def set_bool(self, key, value):
self.store[key] = bool(value)
2023-03-10 03:00:27 +00:00
def get_bool(self, key):
return bool(self.store[key])
def set_int(self, key, value):
self.store[key] = int(value)
def get_int(self, key):
return int(self.store.get(key, 0))
2023-03-10 03:00:27 +00:00
def set_str(self, key, value):
self.store[key] = value
2023-03-10 03:00:27 +00:00
def get_str(self, key):
return str(self.store.get(key, ""))
2023-03-10 04:07:12 +00:00
def set_hash_int(self, ns, key, value):
if ns not in self.store:
self.store[ns] = {}
self.store[ns][key] = int(value)
2023-03-10 03:00:27 +00:00
def get_hash_int(self, ns, key):
d = self.store.get(ns, {})
return int(d.get(key, 0))
def inc_hash_int(self, ns, key):
if ns not in self.store:
self.store[ns] = {}
2023-03-10 03:00:27 +00:00
if key not in self.store[ns]:
self.store[ns][key] = 0
else:
self.store[ns][key] += 1
def dec_hash_int(self, ns, key):
if ns not in self.store:
self.store[ns] = {}
2023-03-10 03:00:27 +00:00
if key not in self.store[ns]:
self.store[ns][key] = 0
else:
self.store[ns][key] -= 1
def get_all_hash_int(self, ns):
2023-03-10 04:07:12 +00:00
if ns in self.store:
return [{str(k): int(v)} for k,v in self.store[ns].items()]
else:
return []
2023-03-10 03:00:27 +00:00
def del_hash(self, ns, key):
del self.store[ns][key]
class RedisStorage(Storage):
def __init__(self, redis_uri):
self.conn = redis.from_url(redis_uri)
self.conn.ping()
2023-03-11 17:06:36 +00:00
def exists(self, key):
return bool(self.conn.exists(key))
2023-03-10 03:00:27 +00:00
def set_bool(self, key, value):
self.conn.set(key, "1" if value else "0")
2023-03-10 03:00:27 +00:00
def get_bool(self, key):
return bool(self.conn.get(key))
def set_int(self, key, value):
self.conn.set(key, str(value))
def get_int(self, key):
v = self.conn.get(key)
if v is None:
return 0
else:
return v
2023-03-10 03:00:27 +00:00
def set_str(self, key, value):
self.conn.set(key, value)
2023-03-10 03:00:27 +00:00
def get_str(self, key):
v = self.conn.get(key)
if v is None:
return ""
else:
return v.decode('utf-8')
2023-03-10 03:00:27 +00:00
def get_hash_int(self, ns, key):
v = self.conn.hget(ns, key)
if v is None:
return 0
else:
return int(v)
2023-03-10 04:07:12 +00:00
def set_hash_int(self, ns, key, value):
self.conn.hset(ns, key, value)
2023-03-10 03:00:27 +00:00
def inc_hash_int(self, ns, key):
return int(self.conn.hincrby(ns, key))
def dec_hash_int(self, ns, key):
return int(self.conn.hincrby(ns, key, -1))
def get_all_hash_int(self, ns):
2023-03-10 04:07:12 +00:00
return {k.decode("utf-8"): int(v) for k,v in self.conn.hgetall(ns).items()}
2023-03-10 03:00:27 +00:00
def del_hash(self, ns, key):
2023-03-10 04:07:12 +00:00
self.conn.hdel(ns, key)
2023-03-10 03:00:27 +00:00
def setup(storage_uri):
global storage
if storage_uri.startswith("memory://"):
storage = MemoryStorage()
elif storage_uri.startswith("redis://"):
storage = RedisStorage(storage_uri)
else:
raise Exception("Invalid storage URI: " + storage_uri)
return storage