diff --git a/tools/datafiles.py b/tools/datafiles.py index d87fc08..c8f2795 100644 --- a/tools/datafiles.py +++ b/tools/datafiles.py @@ -1,5 +1,6 @@ import json import os +import pickle class DataFile(dict): @@ -33,10 +34,24 @@ class JSONFile(DataFile): c = f.read() if len(c) > 0: - json_dict = json.loads(c) - for k in json_dict: - self[k] = json_dict[k] + self.update(json.loads(c)) 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: + self.update(pickle.loads(c)) + + def save(self) -> None: + with open(self.filename, "wb") as f: + pickle.dump(self.copy(), f)