BATCAVE/ffstatus/storage.py

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

import cPickle as pickle
import logging
import os

class Storage:
    data = None

    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_file = open(os.path.join(storage_dir, 'storage.dat'), '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(
                os.path.join(storage_dir, self.storage_file.name),
                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):
        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.data = None