# 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)