Source code for xmm.pipeline.steps.flatten

from .base import PipelineStep


class FlattenStep(PipelineStep):
    """
    Convert nested dicts to flat dicts.

    :param int depth: Number of nested levels to flatten. Set to -1 for infinite depth.

    **Examples**:

    Input::

        {
          'fields': {
            'name': {
              'de': 'Deutsch',
              'en': 'English',
            },
          },
        }

    Output::

        {
          'fields.name.de': 'Deutsch',
          'fields.name.en': 'English',
        }
    """

    def __init__(self, depth=-1):
        self.depth = depth

    def flatten(self, dct, level=0):
        """Flatten a dict."""
        for key, value in dct.items():
            if isinstance(value, dict) and (level < self.depth or self.depth < 0):
                for subkey, subvalue in self.flatten(value, level=level + 1):
                    yield key + '.' + subkey, subvalue
            else:
                yield key, value

    def map(self, state):
        for dataset in state:
            yield dataset.__class__(self.flatten(dataset))

    def process_step(self, state, context):
        state = self.map(state)
        return state, context