CODEWITHSUNDEEP
pythonpython-3.xgeneratorlist-comprehensionpython-itertools
intuited
intuited

Reputation: 24084

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

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

Upvotes: 2

Views: 501

Answers (6)

pylang
pylang

Reputation: 44605

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.

Upvotes: 1

intuited
intuited

Reputation: 24084

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.

Upvotes: 0

Har
Har

Reputation: 3918

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)

Upvotes: 1

Daweo
Daweo

Reputation: 36763

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.

Upvotes: 0

blueteeth
blueteeth

Reputation: 3575

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))

Upvotes: 2

Ch3steR
Ch3steR

Reputation: 20689

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']

Upvotes: 3

Related Questions