from procgraph import Block, ETERNITY
from procgraph.block_utils import make_sure_dir_exists
import pickle


class Pickle(Block):
    ''' Dumps the input as a :py:mod:`pickle` file. '''
    Block.alias('pickle')
    Block.config('file', 'File to write to.')
    Block.input('x', 'Anything pickable.') 
        
    def write(self, x, filename):
        make_sure_dir_exists(filename)
        with open(filename, 'wb') as f:
            pickle.dump(x, f, protocol=pickle.HIGHEST_PROTOCOL)
 
    def update(self):
        self.write(self.input.x, self.config.file + '.part')
        
    def finish(self):
        self.write(self.input.x, self.config.file)


class PickleLoad(Block):
    ''' Dumps the input as a :py:mod:`pickle` file. '''
    Block.alias('pickle_load')
    Block.config('file', 'File to read from.')
    Block.output('x', 'Object read from file') 
        
    def update(self):
        with open(self.config.file, 'rb') as f:
            x = pickle.load(f)
            self.set_output(0, x, timestamp=ETERNITY)
 
 
class PickleGroup(Block):
    ''' Dumps the input as a :py:mod:`pickle` file, in the form
        of a dictionary  signal name -> value.    
    '''
    Block.alias('pickle_group')
    Block.config('file', 'File to write to.')
    Block.input_is_variable('Any number of pickable signals.') 
        
    def write(self, x, filename):
        make_sure_dir_exists(filename)
        with open(filename, 'wb') as f:
            pickle.dump(x, f, protocol=pickle.HIGHEST_PROTOCOL)
 
    def get_dict(self):
        data = {}
        for signal in self.get_input_signals_names():
            data[signal] = self.get_input(signal)
        return data
    
    def update(self):
        self.write(self.get_dict(), self.config.file + '.part')
        
    def finish(self):
        self.write(self.get_dict(), self.config.file)
