# What does “list comprehension” mean? How does it work and how can I use it?

I have the following code:

``````[x ** 2 for x in range(10)]
``````

When I run it in the Python shell, it returns:

``````[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
``````

I’ve searched and it seems this is called a list comprehension, but how does it work?

From the documentation:

List comprehensions provide a concise way to create lists. Common applications are to make new lists where each element is the result of some operations applied to each member of another sequence or iterable, or to create a subsequence of those elements that satisfy a certain condition.

About your question, the list comprehension does the same thing as the following “plain” Python code:

``````>>> l = []
>>> for x in range(10):
...     l.append(x**2)
>>> l
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
``````

How do you write it in one line? Hmm…we can…probably…use `map()` with `lambda`:

``````>>> list(map(lambda x: x**2, range(10)))
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
``````

But isn’t it clearer and simpler to just use a list comprehension?

``````>>> [x**2 for x in range(10)]
[0, 1, 4, 9, 16, 25, 36, 49, 64, 81]
``````

Basically, we can do anything with `x`. Not only `x**2`. For example, run a method of `x`:

``````>>> [x.strip() for x in ('foo\n', 'bar\n', 'baz\n')]
['foo', 'bar', 'baz']
``````

Or use `x` as another function’s argument:

``````>>> [int(x) for x in ('1', '2', '3')]
[1, 2, 3]
``````

We can also, for example, use `x` as the key of a `dict` object. Let’s see:

``````>>> d = {'foo': '10', 'bar': '20', 'baz': '30'}
>>> [d[x] for x in ['foo', 'baz']]
['10', '30']
``````

``````>>> d = {'foo': '10', 'bar': '20', 'baz': '30'}
>>> [int(d[x].rstrip('0')) for x in ['foo', 'baz']]
[1, 3]
``````

And so on.

You can also use `if` or `if...else` in a list comprehension. For example, you only want odd numbers in `range(10)`. You can do:

``````>>> l = []
>>> for x in range(10):
...     if x%2:
...         l.append(x)
>>> l
[1, 3, 5, 7, 9]
``````

Ah that’s too complex. What about the following version?

``````>>> [x for x in range(10) if x%2]
[1, 3, 5, 7, 9]
``````

To use an `if...else` ternary expression, you need put the `if ... else ...` after `x`, not after `range(10)`:

``````>>> [i if i%2 != 0 else None for i in range(10)]
[None, 1, None, 3, None, 5, None, 7, None, 9]
``````

Have you heard about nested list comprehension? You can put two or more `for`s in one list comprehension. For example:

``````>>> [i for x in [[1, 2, 3], [4, 5, 6]] for i in x]
[1, 2, 3, 4, 5, 6]

>>> [j for x in [[[1, 2], ], [[4, 5], ]] for i in x for j in i]
[1, 2, 3, 4, 5, 6]
``````

Let’s talk about the first part, `for x in [[1, 2, 3], [4, 5, 6]]` which gives `[1, 2, 3]` and `[4, 5, 6]`. Then, `for i in x` gives `1`, `2`, `3` and `4`, `5`, `6`.

Warning: You always need put `for x in [[1, 2, 3], [4, 5, 6]]` before `for i in x`:

``````>>> [j for j in x for x in [[1, 2, 3], [4, 5, 6]]]
Traceback (most recent call last):
File "<input>", line 1, in <module>
NameError: name 'x' is not defined
``````

We also have set comprehensions, dict comprehensions, and generator expressions.

set comprehensions and list comprehensions are basically the same, but the former returns a set instead of a list:

``````>>> {x for x in [1, 1, 2, 3, 3, 1]}
{1, 2, 3}
``````

It’s the same as:

``````>>> set([i for i in [1, 1, 2, 3, 3, 1]])
{1, 2, 3}
``````

A dict comprehension looks like a set comprehension, but it uses `{key: value for key, value in ...}` or `{i: i for i in ...}` instead of `{i for i in ...}`.

For example:

``````>>> {i: i**2 for i in range(5)}
{0: 0, 1: 1, 2: 4, 3: 9, 4: 16}
``````

And it equals:

``````>>> d = {}
>>> for i in range(5):
...     d[i] = i**2
>>> d
{0: 0, 1: 1, 2: 4, 3: 9, 4: 16}
``````

Does `(i for i in range(5))` give a tuple? No!, it’s a generator expression. Which returns a generator:

``````>>> (i for i in range(5))
<generator object <genexpr> at 0x7f52703fbca8>
``````

It’s the same as:

``````>>> def gen():
...     for i in range(5):
...         yield i
>>> gen()
<generator object gen at 0x7f5270380db0>
``````

And you can use it as a generator:

``````>>> gen = (i for i in range(5))
>>> next(gen)
0
>>> next(gen)
1
>>> list(gen)
[2, 3, 4]
>>> next(gen)
Traceback (most recent call last):
File "<input>", line 1, in <module>
StopIteration
``````

Note: If you use a list comprehension inside a function, you don’t need the `[]` if that function could loop over a generator. For example, `sum()`:

``````>>> sum(i**2 for i in range(5))
30
``````

Related (about generators): Understanding Generators in Python.

Scroll Up