BATCAVE/ffstatus/storage.py

#!/usr/bin/python
# -*- coding: utf-8 -*-

import cPickle as pickle
import logging
import os

from .basestorage import BaseStorage

class Storage(BaseStorage):
    """Provides file-based persistency for BaseStorage"""

    def __init__(self, storage_dir):
        self.logger = logging.getLogger('Storage')
        self.storage_dir = storage_dir

        self.logger.debug('Using storage at \'{0}\''.format(storage_dir))
        self.storage_filename = os.path.join(storage_dir, 'storage.dat')

        self.storage_file = None

    def open(self):
        self.storage_file = open(self.storage_filename, 'a+b')
        try:
            self.storage_file.seek(0, os.SEEK_SET)
            self.data = pickle.load(self.storage_file)

        except EOFError:
            self.logger.warn('The storage file was empty. I\'ll assume this is the first start with this storage directory.')

        except Exception as err:
            self.logger.error('Loading the storage failed. Please resolve the error and/or remove the storage file \'{0}\': {1}'.format(
                self.storage_filename,
                str(err),
            ))
            raise err

        if self.data is None:
            self.data = {}
        self.logger.info('Opened storage with ' + str(len(self.data)) + ' entries.')

    def save(self):
        if self.storage_file is None:
            return

        self.storage_file.seek(0, os.SEEK_SET)
        self.storage_file.truncate()
        pickle.dump(self.data, self.storage_file, protocol=2)
        self.storage_file.flush()

    def close(self):
        self.save()
        self.storage_file.close()
        self.storage_file = None
        self.data = None