diff --git a/tools/datafiles.py b/tools/datafiles.py index d87fc08..d589441 100644 --- a/tools/datafiles.py +++ b/tools/datafiles.py @@ -1,5 +1,6 @@ import json import os +import pickle class DataFile(dict): @@ -40,3 +41,21 @@ class JSONFile(DataFile): def save(self): with open(self.filename, "wt") as f: f.write(json.dumps(self.copy(), indent=4)) + + +class PickleFile(DataFile): + def __init__(self, filename: str, create: bool) -> None: + super().__init__(filename, create) + + def load(self) -> None: + with open(self.filename, "rb") as f: + c = f.read() + + if len(c) > 0: + pickle_dict = pickle.loads(c) + for k in pickle_dict: + self[k] = pickle_dict[k] + + def save(self) -> None: + with open(self.filename, "wb") as f: + pickle.dump(self.copy(), f)