Sunday, April 10, 2005

python list comprehensions can be intuitive if you squint

Sexy title, eh?

I read a collaborator's code last night and was struck by the number of times you see things like

newList=[]
....for thisThing in thatList:
........for thisOtherThing in thisOtherList:
............newList.append( somefunc( thisThing, thisOtherThing))

and then we can do something with the newList (or dictionaries; similar story).

This lead me to see if I could make list comprehensions be intuitive. I think I succeeded.

To cut to the chase, the above can be rewritten as

[somefunc(thisThing,thatThing) for thisThing in thatList for thisOtherThing in thisOtherList]


Its not more easy to read with long-winded names like this-- (and easy-reading is usually the highest priority for me and python)-- but it is concise and versatile.

So let's try to make them our friend

Given this,
>>> vec = [2, 4, 6]
>>> newList= [x*3 for x in vec]
>>> newList
[6, 12, 18]

vs, the usual way...

newList=[]
for x in vec:
newList.append( 3*x )

To me, [x*3 for x in vec]was never intuitive. (Apparently the syntax comes from mathematical set theory.)

But here's how I can make it livable: give me a list whose elements are 3*x
where each x comes out of vec

[ give me a list
3*x whose elements have the value 3*x
for x in vec for each x in vec ]

So here are some things you can what you can do with this.

Add inner loops:
>>>[x*y for x in vec for y in vec]
[4, 8, 12, 8, 16, 24, 12, 24, 36]

You you can add conditionals. (this is what python's filter function does)
Let's filter extract all the numbers divisible by 3
>>> [x*y for x in vec for y in vec if x*y % 3 == 0]
[12, 24, 12, 24, 36]

Generate lists of tuples or anything else.
Let's see what numbers generated the last set...
[(x , y, x*y) for x in vec for y in vec if x*y % 3 == 0]
[(2, 6, 12), (4, 6, 24), (6, 2, 12), (6, 4, 24), (6, 6, 36)]

Two times 6 is 12, 4 * 6 is 24... oh, right!-- at least one of the factors must be divisible by 3. You can put multiple ifs in here too.

And, as in our first example, invoke other functions
[multiply(x,y) x in vec for y in vec ]
[2, 4, 6 ]

You could probably extend and nest these things hugely, and achieve all the unreadability of lisp.

Don't do that.

Coincidentally, this topic is in the air these days.
See The fate of reduce() in Python 3000 by Guido van Rossum