Tuesday, January 1, 2008

Tablet Python #3 - List Comprehension

A Happy New Year to everyone! In today's episode of Tablet Python I'm going to tell you more about list comprehension which we were already using in the last episode.

Introduction

List comprehension is a unique feature of the Python language (correct me if I'm wrong). It might look intimidating at first, but believe me, it can soon become one of your best friends in Python. It's so elegant and simple!

But let's approach this step by step.

Step 1: You can copy lists

Let's define a list
a = [1, 2, 3]

You can make a copy of a by iterating through the elements in a, like this:
copy_of_a = []
for element in a:
copy_of_a.append(element)

This is about the way you would do it in any other language. But with list comprehension, you would simply write:
copy_of_a = [ element for element in a ]

This simply says, build a list and for every element in a, put element into this list. Then assign the new list to the variable copy_of_a.

For completeness, I want to tell you, however, that the easiest way to copy a list in Python goes by "slicing":copy_of_a = a[:]

Step 2: You can build modified copies

Let's assume you have a list of filesystem paths and would like to strip off the path so that only the filename remains.

This is our list:
a = ["/home/user/file1", "/media/mmc1/file2", "/media/mmc2/file3"]

And this is what we want to get:
["file1", "file2", "file3"]

The classic way goes like this:
b = []
for element in a:
b.append(os.path.basename(element))

But with list comprehension, you'd do it like this:
b = [ os.path.basename(element) for element in a ]

Or in words: build a list, and for every element in a, apply os.path.basename() on element and put the result into the new list. Then assign the new list to the variable b.

Step 3: You can filter lists

Now let's take a look at the list comprehension construct from the last episode:
resources = [ f for f in os.listdir(_RESOURCE_PATH) if f.endswith(".png") ]

Can you guess what it does?
Build a list, and for every f in the result of os.listdir(_RESOURCE_PATH), put f into the new list if f.endswith(".png") returns True. Then assign the new list to the variable resources.
It's the same as:
resources = []
for f in os.listdir(_RESOURCE_PATH):
if f.endswith(".png"):
resources.append(f)

... but way shorter, less error-prone, and easier to read!

Step 4: You can combine modifying and filtering

Let's take this to the extreme by modifying and filtering a list at the same time:
a = [ str(c) for c in range(100) if (c % 5 == 0) ]

This builds a list a which contains as strings those numbers in the range between 0 and 99, which can be divided by 5.
This would be
a = []
for c in range(100):
if (c % 5 == 0):
a.append(str(c))

for all those poor people who cannot use list comprehension.

Step 5: Still confused? It's simply a mathematical set notation!

Do you remember this notation?

List comprehension in Python is just the same!
B = [ x*x for x in A if (x % 5 == 0) ]

And that's the secret key which helps you understand it.

Conclusion

Generally, list comprehension helps you write cleaner code. It expresses how you think about solving the problem instead of describing every step necessary to solve the problem. Once you get used to it, it's a much more intuitive way of working with lists.
A feature of higher-level (functional) languages is that you can make the computer solve problems by describing what you want, instead of giving step-by-step instructions on how to solve it.
However, as with every powerful tool, use it wisely. It's e.g. generally not a good idea to use list comprehension to shorten code with side-effects. Don't use list comprehension to make your code shorter, but to make it understandable.

3 comments:

Unknown said...

Haskell and Erlang also have list comprehensions.

Duncan Booth said...

Slicing isn't the recommended way to copy a list: the recommended way is simply to construct a new list:

copy_of_a = list(a)

This will convert any sequence or iterator to a list, even if it doesn't support slicing. Also the same technique works for other mutable types such as dict (for immutable objects it will convert the type but if the object is already the correct type it doesn't copy it).

ustunozgur said...

Thanks for the article. The final point is vital for understanding comprehension easily.

It seems a number of languages support this feature:

http://en.wikipedia.org/wiki/List_comprehension