astrXbian/.install/.kodi/addons/script.module.xbmcswift2/lib/xbmcswift2/storage.py

'''
    xbmcswift2.storage
    ~~~~~~~~~~~~~~~~~~

    This module contains persistent storage classes.

    :copyright: (c) 2012 by Jonathan Beluch
    :license: GPLv3, see LICENSE for more details.
'''
import os
import csv
import json
import time
try:
    import cPickle as pickle
except ImportError:
    import pickle
import shutil
import collections
from datetime import datetime
from xbmcswift2.logger import log


class _PersistentDictMixin(object):
    ''' Persistent dictionary with an API compatible with shelve and anydbm.

    The dict is kept in memory, so the dictionary operations run as fast as
    a regular dictionary.

    Write to disk is delayed until close or sync (similar to gdbm's fast mode).

    Input file format is automatically discovered.
    Output file format is selectable between pickle, json, and csv.
    All three serialization formats are backed by fast C implementations.
    '''

    def __init__(self, filename, flag='c', mode=None, file_format='pickle'):
        self.flag = flag                    # r=readonly, c=create, or n=new
        self.mode = mode                    # None or an octal triple like 0644
        self.file_format = file_format      # 'csv', 'json', or 'pickle'
        self.filename = filename
        if flag != 'n' and os.access(filename, os.R_OK):
            log.debug('Reading %s storage from disk at "%s"',
                      self.file_format, self.filename)
            fileobj = open(filename, 'rb' if file_format == 'pickle' else 'r')
            with fileobj:
                self.load(fileobj)

    def sync(self):
        '''Write the dict to disk'''
        if self.flag == 'r':
            return
        filename = self.filename
        tempname = filename + '.tmp'
        fileobj = open(tempname, 'wb' if self.file_format == 'pickle' else 'w')
        try:
            self.dump(fileobj)
        except Exception:
            os.remove(tempname)
            raise
        finally:
            fileobj.close()

        # shutil error (SameFileError when performing copyfile)
        if os.path.exists(self.filename):
            os.remove(self.filename)

        shutil.move(tempname, self.filename)    # atomic commit
        if self.mode is not None:
            os.chmod(self.filename, self.mode)

    def close(self):
        '''Calls sync'''
        self.sync()

    def __enter__(self):
        return self

    def __exit__(self, *exc_info):
        self.close()

    def dump(self, fileobj):
        '''Handles the writing of the dict to the file object'''
        if self.file_format == 'csv':
            csv.writer(fileobj).writerows(self.raw_dict().items())
        elif self.file_format == 'json':
            json.dump(self.raw_dict(), fileobj, separators=(',', ':'))
        elif self.file_format == 'pickle':
            pickle.dump(dict(self.raw_dict()), fileobj, 2)
        else:
            raise NotImplementedError('Unknown format: ' +
                                      repr(self.file_format))

    def load(self, fileobj):
        '''Load the dict from the file object'''
        # try formats from most restrictive to least restrictive
        for loader in (pickle.load, json.load, csv.reader):
            fileobj.seek(0)
            try:
                return self.initial_update(loader(fileobj))
            except Exception as e:
                pass
        raise ValueError('File not in a supported format')

    def raw_dict(self):
        '''Returns the underlying dict'''
        raise NotImplementedError


class _Storage(collections.MutableMapping, _PersistentDictMixin):
    '''Storage that acts like a dict but also can persist to disk.

    :param filename: An absolute filepath to reprsent the storage on disk. The
                     storage will loaded from this file if it already exists,
                     otherwise the file will be created.
    :param file_format: 'pickle', 'json' or 'csv'. pickle is the default. Be
                        aware that json and csv have limited support for python
                        objets.

    .. warning:: Currently there are no limitations on the size of the storage.
                 Please be sure to call :meth:`~xbmcswift2._Storage.clear`
                 periodically.
    '''

    def __init__(self, filename, file_format='pickle'):
        '''Acceptable formats are 'csv', 'json' and 'pickle'.'''
        self._items = {}
        _PersistentDictMixin.__init__(self, filename, file_format=file_format)

    def __setitem__(self, key, val):
        self._items.__setitem__(key, val)

    def __getitem__(self, key):
        return self._items.__getitem__(key)

    def __delitem__(self, key):
        self._items.__delitem__(key)

    def __iter__(self):
        return iter(self._items)

    def __len__(self):
        return self._items.__len__

    def raw_dict(self):
        '''Returns the wrapped dict'''
        return self._items

    initial_update = collections.MutableMapping.update

    def clear(self):
        super(_Storage, self).clear()
        self.sync()


class TimedStorage(_Storage):
    '''A dict with the ability to persist to disk and TTL for items.'''

    def __init__(self, filename, file_format='pickle', TTL=None):
        '''TTL if provided should be a datetime.timedelta. Any entries
        older than the provided TTL will be removed upon load and upon item
        access.
        '''
        self.TTL = TTL
        _Storage.__init__(self, filename, file_format=file_format)

    def __setitem__(self, key, val, raw=False):
        if raw:
            self._items[key] = val
        else:
            self._items[key] = (val, time.time())

    def __getitem__(self, key):
        val, timestamp = self._items[key]
        if self.TTL and (datetime.utcnow() -
            datetime.utcfromtimestamp(timestamp) > self.TTL):
            del self._items[key]
            return self._items[key][0]  # Will raise KeyError
        return val

    def initial_update(self, mapping):
        '''Initially fills the underlying dictionary with keys, values and
        timestamps.
        '''
        for key, val in mapping.items():
            _, timestamp = val
            if not self.TTL or (datetime.utcnow() -
                datetime.utcfromtimestamp(timestamp) < self.TTL):
                self.__setitem__(key, val, raw=True)