Source code for bonobo.nodes.io.pickle

import pickle

from bonobo.config import Option, use_context
from bonobo.constants import NOT_MODIFIED
from bonobo.nodes.io.base import FileHandler
from bonobo.nodes.io.file import FileReader, FileWriter


class PickleHandler(FileHandler):
    """

    .. attribute:: item_names

        The names of the items in the pickle, if it is not defined in the first item of the pickle.

    """

    fields = Option(tuple, required=False)


[docs]@use_context class PickleReader(FileReader, PickleHandler): """ Reads a Python pickle object and yields the items in dicts. """ mode = Option(str, default='rb')
[docs] def read(self, file, context, *, fs): data = pickle.load(file) # if the data is not iterable, then wrap the object in a list so it may be iterated if isinstance(data, dict): is_dict = True iterator = iter(data.items()) else: is_dict = False try: iterator = iter(data) except TypeError: iterator = iter([data]) if not context.output_type: context.set_output_fields(self.fields or next(iterator)) fields = context.get_output_fields() fields_length = len(fields) for row in iterator: if len(row) != fields_length: raise ValueError('Received an object with {} items, expected {}.'.format(len(row), fields_length))
yield tuple(row.values() if is_dict else row)
__call__ = read
[docs]@use_context class PickleWriter(FileWriter, PickleHandler): mode = Option(str, default='wb')
[docs] def write(self, file, context, item, *, fs): """ Write a pickled item to the opened file. """ context.setdefault('lineno', 0) file.write(pickle.dumps(item)) context.lineno += 1
return NOT_MODIFIED
__call__ = write