multiple variables in list comprehension?

Question:

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()])
Asked By: Dave

||

Answers:

If you really want to do it in one line, you can do something like this, although it’s not the clearest code (first line is the code, second line is the output).

>>> [(lambda x, y, z: [int(x), y.strip(), z.strip()])(*s.split(",")) for s in inputs]
[[1, 'foo', 'bar'], [2, 'tom', 'jerry']]

Or this.

>>> [(lambda x: [int(x[0]), x[1].strip(), x[2].strip()])(s.split(",")) for s in inputs]
[[1, 'foo', 'bar'], [2, 'tom', 'jerry']

Edit:
See jonrsharpe’s comment for the best answer IMHO.

Answered By: Rok Povsic

You can use map with a list comprehension

def clean(x):
   return [int(x[0]), x[1].strip(), x[2].strip()]

map(clean,[s.split(",") for s in inputs])
# Output: [[1, 'foo', 'bar'], [2, 'tom', 'jerry']]

with a lambda function:

map(lambda x: [int(x[0]), x[1].strip(), x[2].strip()],[s.split(",") for s in inputs])
Answered By: user2314737

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 make use of it here.

Answered By: kindall

Another possible solution

>>> inputs = [[1, "foo", "bar"], [2, "tom", "jerry"]]
>>> list_of_dicts = [{"var{}".format(k): v for k, v in enumerate(s, start=1)} for s in inputs]
>>> list_of_dicts[0]['var1']
1
>>> list_of_dicts[0]['var2']
'foo'
Answered By: shaktimaan

Thanks for all the suggestions – it’s great to see the variety of possible techniques, as well as a counterexample to the zen of Python’s “There should be one — and preferably only one — obvious way to do it.”

All 4 solutions are equally beautiful, so it’s a bit unfair to have to give the coveted green check to only one of them. I agree with the recommendations that #1 is the cleanest and best approach. #2 is also straightforward to understand, but having to use a lambda inside a comprehension seems a bit off. #3 is nice in creating an iterator with map but gets a tiny demerit for needing the extra step of iterating over it. #4 is cool for pointing out that nested for’s are possible — if I can remember that they go in “first, second” order instead of “inner, outer”. Since #1 is not in the form of an answer, #4 gets the check for most surprising.

Thanks again to all.

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

outputs1 = [[int(x), y.strip(), z.strip()] for x,y,z in (s.split(',') for s in inputs)]
print("1:", outputs1)       # jonrsharpe

outputs2 = [(lambda x, y, z: [int(x), y.strip(), z.strip()])(*s.split(",")) for s in inputs]
print("2:", outputs2)       # yper

outputs3 = [z for z in map(lambda x: [int(x[0]), x[1].strip(), x[2].strip()],[s.split(",") for s in inputs])]
print("3:", outputs3)       # user2314737

outputs4 = [[int(x), y.strip(), z.strip()] for s in inputs for (x, y, z) in [s.split(",")]]
print("4:", outputs4)       # kindall

Results:

1: [[1, 'foo', 'bar'], [2, 'tom', 'jerry']]
2: [[1, 'foo', 'bar'], [2, 'tom', 'jerry']]
3: [[1, 'foo', 'bar'], [2, 'tom', 'jerry']]
4: [[1, 'foo', 'bar'], [2, 'tom', 'jerry']]
Answered By: Dave
Categories: questions Tags: ,
Answers are sorted by their score. The answer accepted by the question owner as the best is marked with
at the top-right corner.