1
0
mirror of https://github.com/LibreTranslate/LibreTranslate.git synced 2024-11-05 07:20:13 +01:00
LibreTranslate/libretranslate/storage.py
2023-03-09 22:00:27 -05:00

147 lines
3.8 KiB
Python

import redis
storage = None
def get_storage():
return storage
class Storage:
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")
def set_hash_value(self, ns, key, value):
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")
def get_hash_keys(self, ns):
raise Exception("not implemented")
def del_hash(self, ns, key):
raise Exception("not implemented")
class MemoryStorage(Storage):
def __init__(self):
self.store = {}
def set_bool(self, key, value):
self.store[key] = bool(value)
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))
def set_str(self, key, value):
self.store[key] = value
def get_str(self, key):
return str(self.store.get(key, ""))
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] = {}
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] = {}
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):
return [{str(k): int(v)} for k,v in self.store[ns].items()]
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()
def set_bool(self, key, value):
self.conn.set(key, "1" if value else "0")
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
def set_str(self, key, value):
self.conn.set(key, value)
def get_str(self, key):
v = self.conn.get(key)
if v is None:
return ""
else:
return v.decode('utf-8')
def get_hash_int(self, ns, key):
v = self.conn.hget(ns, key)
if v is None:
return 0
else:
return int(v)
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):
return [{k.decode("utf-8"): int(v)} for k,v in self.conn.hgetall(ns).items()]
def del_hash(self, ns, key):
conn.hdel(ns, key)
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