PyLink/structures.py

"""
structures.py - PyLink data structures module.

This module contains custom data structures that may be useful in various situations.
"""

import collections
import json

class KeyedDefaultdict(collections.defaultdict):
    """
    Subclass of defaultdict allowing the key to be passed to the default factory.
    """
    def __missing__(self, key):
        if self.default_factory is None:
            # If there is no default factory, just let defaultdict handle it
            super().__missing__(self, key)
        else:
            value = self[key] = self.default_factory(key)
            return value

class DataStore:
    # will come into play with subclassing and db version upgrading
    initial_version = 1

    def __init__(self, name, filename, db_format='json', save_frequency={'seconds': 30}):
        self.name = name

        self._filename = os.path.abspath(os.path.expanduser(filename))
        self._tmp_filename = self._filename + '.tmp'
        log.debug('(db:{}) database path set to {}'.format(self.name, self._filename))

        self._format = db_format
        log.debug('(db:{}) format set to {}'.format(self.name, self._format))

        self._save_frequency = timedelta(**save_frequency).total_seconds()
        log.debug('(db:{}) saving every {} seconds'.format(self.name, self._save_frequency))

    def create_or_load(self):
        log.debug('(db:{}) creating/loading datastore using {}'.format(self.name, self._format))

        if self._format == 'json':
            self._store = {}
            self._store_lock = threading.Lock()

            log.debug('(db:{}) loading json data store from {}'.format(self.name, self._filename))
            try:
                self._store = json.loads(open(self._filename, 'r').read())
            except (ValueError, IOError, FileNotFoundError):
                log.exception('(db:{}) failed to load existing db, creating new one in memory'.format(self.name))
                self.put('db.version', self.initial_version)
        else:
            raise Exception('(db:{}) Data store format [{}] not recognised'.format(self.name, self._format))

    def save_callback(self, starting=False):
        """Start the DB save loop."""
        if self._format == 'json':
            # don't actually save the first time
            if not starting:
                self.save()

            # schedule saving in a loop.
            self.exportdb_timer = threading.Timer(self._save_frequency, self.save_callback)
            self.exportdb_timer.name = 'PyLink {} save_callback Loop'.format(self.name)
            self.exportdb_timer.start()
        else:
            raise Exception('(db:{}) Data store format [{}] not recognised'.format(self.name, self._format))

    def save(self):
        log.debug('(db:{}) saving datastore'.format(self.name))
        if self._format == 'json':
            with open(self._tmp_filename, 'w') as store_file:
                store_file.write(json.dumps(self._store))
            os.rename(self._tmp_filename, self._filename)

    # single keys
    def __contains__(self, key):
        if self._format == 'json':
            return key in self._store

    def get(self, key, default=None):
        if self._format == 'json':
            return self._store.get(key, default)

    def put(self, key, value):
        if self._format == 'json':
            # make sure we can serialize the given data
            # so we don't choke later on saving the db out
            json.dumps(value)

            self._store[key] = value

            return True

    def delete(self, key):
        if self._format == 'json':
            try:
                with self._store_lock:
                    del self._store[key]
            except KeyError:
                # key is already gone, nothing to do
                ...

            return True

    # multiple keys
    def list_keys(self, prefix=None):
        """Return all key names. If prefix given, return only keys that start with it."""
        if self._format == 'json':
            keys = []

            with self._store_lock:
                for key in self._store:
                    if prefix is None or key.startswith(prefix):
                        keys.append(key)

            return keys

    def delete_keys(self, prefix):
        """Delete all keys with the given prefix."""
        if self._format == 'json':
            with self._store_lock:
                for key in tuple(self._store):
                    if key.startswith(prefix):
                        del self._store[key]