84 lines
2.1 KiB
Python
84 lines
2.1 KiB
Python
# Python imports
|
|
import os
|
|
import pickle
|
|
from typing import Dict, TypeVar, Generic, Optional
|
|
from distutils.dir_util import mkpath
|
|
|
|
# Lib imports
|
|
|
|
# Application imports
|
|
|
|
|
|
|
|
|
|
Key = TypeVar('Key')
|
|
Value = TypeVar('Value')
|
|
Records = Dict[Key, Value]
|
|
|
|
|
|
class KeyValueDb(Generic[Key, Value]):
|
|
"""
|
|
Key-value in-memory database
|
|
Use open() method to load DB from a file and commit() to save it
|
|
"""
|
|
|
|
_name = None # type: str
|
|
_records = None # type: Records
|
|
|
|
def __init__(self, basename: str):
|
|
"""
|
|
:param str basename: path to db file
|
|
"""
|
|
self._name = basename
|
|
self.set_records({})
|
|
|
|
def open(self) -> 'KeyValueDb':
|
|
""" Create a new data base or open existing one... """
|
|
if os.path.exists(self._name):
|
|
if not os.path.isfile(self._name):
|
|
raise IOError("%s exists and is not a file" % self._name)
|
|
|
|
if os.path.getsize(self._name) == 0:
|
|
return self
|
|
|
|
with open(self._name, 'rb') as _in: # binary mode
|
|
self.set_records(pickle.load(_in))
|
|
else:
|
|
mkpath(os.path.dirname(self._name))
|
|
self.commit()
|
|
|
|
return self
|
|
|
|
def commit(self) -> 'KeyValueDb':
|
|
""" Write the database to a file... """
|
|
with open(self._name, 'wb') as out:
|
|
pickle_protocol = 2 # use 2 for BC with Ulauncher 4
|
|
pickle.dump(self._records, out, pickle_protocol)
|
|
out.close()
|
|
|
|
return self
|
|
|
|
def remove(self, key: Key) -> bool:
|
|
"""
|
|
:param str key:
|
|
:type: bool
|
|
:return: True if record was removed
|
|
"""
|
|
try:
|
|
del self._records[key]
|
|
return True
|
|
except KeyError:
|
|
return False
|
|
|
|
def set_records(self, records: Records):
|
|
self._records = records
|
|
|
|
def get_records(self) -> Records:
|
|
return self._records
|
|
|
|
def put(self, key: Key, value: Value) -> None:
|
|
self._records[key] = value
|
|
|
|
def find(self, key: Key, default: Value = None) -> Optional[Value]:
|
|
return self._records.get(key, default)
|