Split and flatten a list of Strings and None values using a comprehension

user_192812 0
Viewed 238

Given a list which contains both strings and None values, in which some of the strings have embedded newlines, I wish to split the strings with newlines into multiple strings and return a flattened list.

I've written code to do this using a generator function, but the code is rather bulky and I'm wondering if it's possible to do it more concisely using a list comprehension or a function from the itertools module. itertools.chain doesn't seem to be able to decline to iterate any non-iterable elements.

def expand_newlines(lines):
    r"""Split strings with newlines into multiple strings.

    >>> l = ["1\n2\n3", None, "4\n5\n6"]
    >>> list(expand_newlines(l))
    ['1', '2', '3', None, '4', '5', '6']
    """
    for line in lines:
        if line is None:
            yield line
        else:
            for l in line.split('\n'):
                yield l
itertools generator list-comprehension python-3.x python
6 Answers

You can use yield from.

def expand(lines):
    for line in lines:
        if isinstance(line,str):
            yield from line.split('\n')
        elif line is None:
            yield line

list(expand(l))
#['1', '2', '3', None, '4', '5', '6']
user_124164530

Here's a single line, but I think @Ch3steR's solution is more readable.

from itertools import chain

list(chain.from_iterable(i.splitlines() if i is not None and '\n' in i else [i] 
                         for i in lines))
user_19041460

You could use itertools.chain if you did the following

import itertools

def expand_newlines(lines):

    return itertools.chain.from_iterable(x.split("\n") if x else [None]
                                         for x in lines)
user_19324050

Using more_itertools.collapse to flatten nested lists:

Given

import more_itertools as mit


lst = ["1\n2\n3", None, "7\n8\n9"]

Demo

list(mit.collapse([x.split("\n") if x else x for x in lst ]))

# ['1', '2', '3', None, '7', '8', '9']

more_itertools is a third-party package. Install via > pip install more_itertools.

user_45312700

If you might modify list inplace then you might do:

lst = ["1\n2\n3", None, "4\n5\n6"]
for i in range(len(lst))[::-1]:
    if isinstance(lst[i], str):
        lst[i:i+1] = lst[i].split("\n")
print(lst)  # ['1', '2', '3', None, '4', '5', '6']

this solution utilize fact that you might not only get python's list slices, but also assign to them. It moves from right to left, as otherwise I would need to keep count of additional items, which would make it harder.

user_107859750

Similar to @blueteeth's answer but more concise by way of inverting the logic:

import itertools
chainfi = itertools.chain.from_iterable

def expand_newlines(lines):
    r"""Split strings with newlines into multiple strings.

    >>> l = ["1\n2\n3", None, "4\n5\n6"]
    >>> list(expand_newlines(l))
    ['1', '2', '3', None, '4', '5', '6']
    """
    return chainfi([None] if l is None else l.split('\n') for l in lines)

None is the special case so that's what we should be checking for.

This is concise enough that I wouldn't even bother writing a function for it—I just kept it in the function to confirm it works via doctest.

user_1928120
Related