Flatten

flattr.py

__all__ = ('flatten',)

from collections import Iterable

sentinel = []


def flatten(to_flatten):
    flat = _flatten(to_flatten, lambda x: x)
    while iterable(flat) and flat and callable(flat[0]):
        # Iterate through flat. It will consist of pair tuples
        # containing the next continuation to call as well as the value
        # that has been flattened. If the value is anything other than
        # the sentinel, yield it.
        if flat[1] != sentinel:
            yield flat[1]

        # Replace flat with the result of the next continuation,
        # providing the sentinel. When this value is returned, the
        # current branch of continuations has been exhausted.
        flat = flat[0](sentinel)

    # Once flat is no longer iterable, yield its last remaining value.
    yield flat


def _flatten(to_flatten, f, value_or_sentinel=sentinel):
    if not iterable(to_flatten):
        return f(to_flatten)
    if not to_flatten:
        return f(value_or_sentinel)
    l_y = lambda y: _flatten(to_flatten[1:], f, y)
    return lambda x: _flatten(to_flatten[0], l_y, x), value_or_sentinel


def iterable(thing):
    return isinstance(thing, Iterable)


if __name__ == '__main__':
    assert list(flatten([0])) == [0]
    assert list(flatten([0, 1])) == [0, 1]
    assert list(flatten([0, 1, 2])) == [0, 1, 2]

    assert list(flatten((0,))) == [0]
    assert list(flatten((0, 1))) == [0, 1]
    assert list(flatten((0, 1, 2))) == [0, 1, 2]

    assert list(flatten([0, [1]])) == [0, 1]
    assert list(flatten([0, [1, 2]])) == [0, 1, 2]
    assert list(flatten([0, [1, [2]]])) == [0, 1, 2]
    assert list(flatten([0, [1], 2])) == [0, 1, 2]
    assert list(flatten([[0], 1, [2]])) == [0, 1, 2]
    assert list(flatten([[0, 1], 2])) == [0, 1, 2]
    assert list(flatten([[[0, 1, 2]]])) == [0, 1, 2]

    assert list(flatten((0, (1,)))) == [0, 1]
    assert list(flatten((0, (1, 2)))) == [0, 1, 2]
    assert list(flatten((0, (1, (2,))))) == [0, 1, 2]
    assert list(flatten((0, (1,), 2))) == [0, 1, 2]
    assert list(flatten(((0,), 1, (2,)))) == [0, 1, 2]
    assert list(flatten(((0, 1), 2))) == [0, 1, 2]
    assert list(flatten((((0, 1, 2))))) == [0, 1, 2]

    import sys
    n = sys.getrecursionlimit()
    l = None
    for x in range(n * 2):
        if l:
            l = [l, x]
        else:
            l = [x]
    assert list(flatten(l)) == list(range(n * 2))