Dave Dave - 1 month ago 12
Python Question

multiple variables in list comprehension?

I want to create a list of lists from a list of multi-field strings and wonder if it is possible to do so in a comprehension.

Input:

inputs = ["1, foo, bar", "2, tom, jerry"]


Desired output:

[[1, "foo", "bar"], [2, "tom", "jerry"]]


Splitting the string in a comprehension is easy:

>>> [s.split(",") for s in inputs]
[['1', ' foo', ' bar'], ['2', ' tom', ' jerry']]


But I'm having trouble figuring out how to access the columns after the string has been split inside the comprehension, because it would seem to require a variable assignment. The following are not valid Python, but illustrate what I'm looking for:

[[int(x), y.strip(), z.strip() for x,y,z = s.split(",")] for s in inputs]
or
[[int(v[0]), v[1].strip(), v[2].strip() for v = s.split(",")] for s in inputs]


Is there a way to assign variables inside a comprehension so that the output can be composed of functions of the variables? A loop is trivial, but generating a list by transforming inputs sure seems like a "comprehension-ish" task.

outputs = []
for s in inputs:
x,y,z = s.split(",")
outputs.append([int(x), y.strip(), z.strip()])

Answer

You can do this with two for clauses in your list comprehension. The first iterates over the items in the list. The second iterates over a single-item list containing the list derived from splitting the string (which is needed so we can unpack this into three separate variables).

[[int(x), y.strip(), z.strip()] for s in inputs for (x, y, z) in [s.split(",")]]

The for clauses go in a somewhat counterintuitive order, but it matches the way you'd write it as nested for loops.

Jon Sharpe's use of a nested comprehension (generator expression, actually) is similar and probably clearer. The use of multiple for clauses always seems confusing to me; mainly I wanted to see if I could do make use of it here.