import pickle


class Cache:
    def __init__(self, model_version) -> None:
        path_compatible_model_version = model_version.replace("/", "-")
        self.cache_file_name = f"cache_{path_compatible_model_version}.pkl"
        try:
            self.cache = pickle.load(open(self.cache_file_name, "rb"))
        except FileNotFoundError:
            self.cache = {}

    def get(self, key):
        return self.cache.get(key)

    def set(self, key, value):
        self.cache[key] = value

    def save(self):
        pickle.dump(self.cache, open(self.cache_file_name, "wb"))
