Citrusbyte - Array flatten function

flatten.py

def recursive_flatten(nested):
    """Return a flatten array from an arbitrarily nested array of integers

    This function is recursive so it could not be used with lists with high
    levels of nesting to avoid recursion errors

    >>> recursive_flatten([])
    []
    >>> recursive_flatten([1, 2, 3])
    [1, 2, 3]
    >>> recursive_flatten([[1], 2])
    [1, 2]
    >>> recursive_flatten([[1], [2], [3]])
    [1, 2, 3]
    >>> recursive_flatten([[1], [2, 3], [4, [5]]])
    [1, 2, 3, 4, 5]
    >>> a = []
    >>> for i in xrange(1000000):
    ...     a = [a, i]
    >>> # Test list with highly levels of nesting
    >>> recursive_flatten(a) == list(xrange(1000000))    
    True
    >>> 
    """
    flattened = []
    for item in nested:
        if isinstance(item, list):
            flattened.extend(flatten(item))
        else:
            flattened.append(item)
    return flattened


def flatten(nested):
    """Return a flatten array from an arbitrarily nested array of integers

    >>> flatten([])
    []
    >>> flatten([1, 2, 3])
    [1, 2, 3]
    >>> flatten([[1], 2])
    [1, 2]
    >>> flatten([[[[], 1], 2], 3])
    [1, 2, 3]
    >>> flatten([[1], [2], [3]])
    [1, 2, 3]
    >>> flatten([[1], [2, 3], [4, [5]]])
    [1, 2, 3, 4, 5]
    >>> a = []
    >>> for i in xrange(1000000):
    ...     a = [a, i]
    >>> # Test list with highly levels of nesting
    >>> flatten(a) == list(xrange(1000000))    
    True
    >>> 
    """
    # Create a copy of the list to avoid modifying the original parameter
    flattened = nested[:]

    idx = 0
    output = []
    while idx < len(flattened):
        item = flattened[idx]
        if isinstance(item, list):
            if not item:
                # Ignore empty lists
                flattened.pop(idx)
            else:
                # Replace actual item with expanded version of the current item
                flattened[idx:idx + 1] = item

            # Only update index if actual item is not a list
            if not isinstance(flattened[idx], list):
                idx += 1
        else:
            # if it is a number, move on to the next item in list
            idx += 1
    return flattened


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