flatten.py

from collections import Iterable, Mapping
from operator import methodcaller


def flatten(it, map_iter='values'):
    """
    Iterator that walk through nested iterable.
    If the initial item is not iterable, only itself will be yielded.

    Args:
        item     (Object): Any object.
        map_iter    (str): the method to be called when handling mapping

    >>> from collections import OrderedDict
    >>> kwargs = OrderedDict()
    >>> kwargs['file'] = 'test.txt'
    >>> kwargs['data'] = [[0, 1], [-1, 0]]
    >>> tuple(flatten(kwargs))
    ('test.txt', 0, 1, -1, 0)

    >>> tuple(flatten('apple'))
    ('apple',)
    """
    if isinstance(it, str):
        yield it
    elif isinstance(it, Mapping):
        for item in methodcaller(map_iter)(it):
            yield from flatten(item, map_iter=map_iter)
    elif isinstance(it, Iterable):
        for item in it:
            yield from flatten(item, map_iter=map_iter)
    else:
        yield it


def named_flatten(it, key=()):
    """
    Iterator that walk through nested iterable and create tuple key.
    For non Mapping iterable, the key are generated by enumerate.
    If the initial item is not iterable, only itself will be yielded.

    Args:
        item         (Object): Any object.

    >>> from collections import OrderedDict
    >>> kwargs = OrderedDict()
    >>> kwargs['file'] = 'test.txt'
    >>> kwargs['data'] = [[0, 1], [-1, 0]]
    >>> dict(named_flatten(kwargs))['data', 0, 1]
    1

    >>> dict(named_flatten('apple'))[()]
    'apple'
    """
    if isinstance(it, str):
        yield key, it
    elif isinstance(it, Mapping):
        for k, v in it.items():
            yield from named_flatten(v, key=key+(k,))
    elif isinstance(it, Iterable):
        for k, v in enumerate(it):
            yield from named_flatten(v, key=key+(k,))
    else:
        yield key, it


if __name__ == '__main__':
    import doctest
    doctest.testmod()