How to create a listcomp for listing non-list values and then flatten everything using itertools?

Question

I'm asking you here about the issue in the title because I'm both new to Python and especially creating listcomps.

I would like to create a simple function to:

1. check if all args in *args are lists
2. if no then it changes a value into a list
3. flatten everything using itertools

Below you could find my code:

import itertools


def lstfunc(*args):
    lstargs = list(args)
    for i in lstargs:
        if not isinstance(i, list):
            lstargs[lstargs.index(i)] = [i]
    lstargs = list(itertools.chain.from_iterable(lstargs))
    return lstargs


print(lstfunc([1, 2, 3], 4, 5))

Output: [1, 2, 3, 4, 5] (correct).

My question is how to create a listcomp from this part:

    for i in lstargs:
        if not isinstance(i, list):
            lstargs[lstargs.index(i)] = [i]

I tried:

import itertools


def lstfunc(*args):
    lstargs = list(args)
    lstargs = [[i] for i in lstargs if not isinstance(i, list)]
    lstargs = list(itertools.chain.from_iterable(lstargs))
    return lstargs


print(lstfunc([1, 2, 3], 4, 5))

But sadly this retuns only [4, 5].

What's my mistake?

Answer 1

After much pain I found in fact the most simple and pythonic solution (I guess)! I needed the function above to create a 1D list to pass it to another function. However it's not necessary. I just needed an unpacking operator:

def testfunc(*args):
    print(args)

testfunc(*[1, 2, 3], 4)

Output: (1, 2, 3, 4).

This solution automatically flattens all values into a 1D pure iterable (tuple I guess).

Answer 2

The listcomp you want is:

lstargs = [x if isinstance(x, list) else [x] for x in lstargs]

to reproduce the whole list. Your attempt just filtered out the lists, it didn't choose between wrapping and non-wrapping and produce both. You don't need to preconvert args to a list to be clear, so you could simplify the function to just:

def lstfunc(*args):
    # Save eager conversion to temp list by using genexpr instead of listcomp
    lstargs = (x if isinstance(x, list) else [x] for x in args)  # 
    return list(itertools.chain.from_iterable(lstargs))