smart_saver_loader.py

import pickle
import atexit

from threading import Thread, Condition, Lock


class PerlineSaver(object):
    def __init__(self, save_file, index_file):
        self.save_file = save_file
        self.index_file = index_file

        self.seeker = {}
        self.index = {}
        self.no_saved = 0

        self.to_save = []

        self.saver_thread = Thread(target=self._saving_loop, daemon=True)
        self.saver_condition = Condition()
        self.to_save_lock = Lock()
        self.saver_lock = Lock()

        self.saver_thread.start()

        atexit.register(self.cleanup)

    def save(self, i, item):
        with self.to_save_lock, self.saver_condition:
            self.to_save.append((i, item))
            self.saver_condition.notify()

    def _saving_loop(self):
        while True:
            with self.saver_condition:
                self.saver_condition.wait_for(lambda: len(self.to_save) > 0)
            self._pop_save()

    def _pop_save(self):
        with self.to_save_lock:
            i, item = self.to_save.pop(0)
        pickled = pickle.dumps(item)

        with self.saver_lock:
            with open(self.save_file, "ab") as f:
                start = f.tell()
                f.write(pickled + b'\n')
            self.index[i] = {"line_id": self.no_saved, "seek": start}
            self.no_saved += 1

    def cleanup(self):
        print(f"Cleaning saver. Saves all unsaved items ({len(self.to_save)}) and the index...", end="")
        while len(self.to_save) > 0:
            self._pop_save()

        with open(self.index_file, "wb") as f:
            pickle.dump(self.index, f)

        print("Clean up is done")


class PerlineLoader(object):
    def __init__(self, save_file, index_file):
        self.save_file = save_file
        self.index_file = index_file

        with open(index_file, "rb") as f:
            self.index = pickle.load(f)

        self.data = {}

    def load(self, i):
        if i not in self.data.keys():
            data = self._load(i)
            self.data[i] = data
        return self.data[i]

    def _load(self, i):
        to_seek = self.index[i]["seek"]
        with open(self.save_file, "rb") as f:
            f.seek(to_seek)
            line = f.readline()
        item = pickle.loads(line)
        return item


if __name__ == "__main__":
    # Example
    import time
    saver = PerlineSaver("save.test", "index.test")
    for i in range(100):
        saver.save(i, {i: f"item {i}"})
        if len(saver.to_save) > 3:
            print("If can buffer :", saver.to_save[:5])
    while len(saver.to_save) > 0:
        print("But it clears up during time :", len(saver.to_save))
        time.sleep(0.1)

    print("Now you can load")
    loader = PerlineLoader("save.test", "index.test")
    item = loader.load(42)
    print("Item 42 was loaded :", item)