CyberAlphaWolf

import pickle
import logging
import traceback


class PickleDatabase(object):
    """This object is a wrapper to access and handle our
    local pickle file. We'll use this pickle file as a sort of
    database for storing schedules and other info for the agent."""

    def __init__(self, file_path):
        self.file_path = file_path
        try:
            self.data = pickle.load(open(file_path, "rb"))
        except Exception:
            # If something went wrong, the data is either corrupted or missing.
            log = logging.getLogger(self.__class__.__name__)
            log.warn(
                "Unable to open database {}, creating an empty one".format(
                    self.file_path
                )
            )
            self.data = {}
            open(self.file_path, "wb")  # Create the DB file if it doesn't exist
            log.info("Database created: %s" % self.file_path)

    def save(self):
        # Dump current contents of our dict into the pickle file
        pickle.dump(self.data, open(self.file_path, "wb"))

    def __getitem__(self, key):
        return self.data[key]

    def __setitem__(self, key, value):
        self.data[key] = value
        # Save our newly updated dictionary
        self.save()

    def __repr__(self):
        return str(self.data)

    def __contains__(self, key):
        return key in self.data