PythonMastery
reference 5 min read · lesson 8 of 25 in Reference

Python Comprehensions Cheat Sheet

1 · The lesson

read

A comprehension is a for loop, written in one line, that builds a new list, set or dict. Read [n * 2 for n in nums if n > 0] as "for each n in nums, if n > 0, keep n * 2". This page shows every form with examples you can run. For the full tutorial, see Comprehensions.


1. The four forms

The brackets decide what you get back. Round brackets do not make a tuple: they make a generator, an object that makes the values one at a time when something asks for them, without building a list.

FormWritten asGives you
List[expr for x in items]a list
Set{expr for x in items}a set (duplicates dropped)
Dict{key: value for x in items}a dict
Generator(expr for x in items)a generator
python
nums = [3, -1, 4, 0, -2]
[n * n for n in nums]              # [9, 1, 16, 0, 4]
{n % 2 for n in nums}              # {0, 1}
{n: n * n for n in nums[:2]}       # {3: 9, -1: 1}
sum(n * n for n in nums)           # 30  (a generator: no list is built)
tuple(n * 2 for n in nums[:3])     # (6, -2, 8)  (to get a tuple, pass a generator to tuple())

2. It is a loop, in one line

These two pieces of code give the same result. The for and the if come in the same order as in the loop.

python
result = []
for n in nums:
    if n > 0:
        result.append(n * 2)
result                             # [6, 8]
[n * 2 for n in nums if n > 0]     # [6, 8]
+ setup added so this can run · re-runs 1 earlier block on this page first
nums = [3, -1, 4, -1, 5]

# Earlier blocks on this page that this one builds on, re-run quietly first.
_pm_earlier = [
    'nums = [3, -1, 4, 0, -2]\n[n * n for n in nums]              # [9, 1, 16, 0, 4]\n{n % 2 for n in nums}              # {0, 1}\n{n: n * n for n in nums[:2]}       # {3: 9, -1: 1}\nsum(n * n for n in nums)           # 30  (a generator: no list is built)\ntuple(n * 2 for n in nums[:3])     # (6, -2, 8)  (to get a tuple, pass a generator to tuple())',
]
def _pm_run_earlier(_srcs, _g):
    import sys as _s, io as _io, warnings as _w
    for _src in _srcs:
        _o, _e = _s.stdout, _s.stderr
        _s.stdout = _s.stderr = _io.StringIO()
        try:
            with _w.catch_warnings():
                _w.simplefilter('ignore')
                exec(compile(_src, '<earlier block>', 'exec'), _g)
        except (Exception, SystemExit):
            pass
        finally:
            _s.stdout, _s.stderr = _o, _e
    if 'matplotlib.pyplot' in _s.modules:
        _s.modules['matplotlib.pyplot'].close('all')   # their plots belong to their own blocks
_pm_run_earlier(_pm_earlier, globals())
del _pm_earlier, _pm_run_earlier

3. Filter: if at the end

An if after the for is a filter: it decides which items to keep.

python
[n for n in nums if n > 0]               # [3, 4]
[n for n in nums if n % 2 == 0]          # [4, 0, -2]
[n for n in nums if n > 0 if n % 2]      # [3]  (two ifs: both must be true, like and)
[y for n in nums if (y := n * n) > 5]    # [9, 16]
+ setup added so this can run · defines y, nums
# Lightweight mock for objects whose attributes/methods aren't critical
class _AutoMock:
    def __init__(self, name='mock'): self._name = name
    def __getattr__(self, k): return _AutoMock(self._name + '.' + k)
    def __call__(self, *a, **kw):
        print('-> ' + self._name + '() called')
        return _AutoMock(self._name + '()')
    def __repr__(self): return '<mock ' + self._name + '>'
    def __str__(self): return '<mock ' + self._name + '>'
    def __bool__(self): return True
    def __iter__(self): return iter([])
    def __len__(self): return 0
    def __getitem__(self, k): return _AutoMock(self._name + '[...]')
    def __setitem__(self, k, v): pass
    def __enter__(self): return self
    def __exit__(self, *a): return False
    async def __aenter__(self): return self
    async def __aexit__(self, *a): return False
    def __add__(self, o): return self
    def __radd__(self, o): return self
    def __sub__(self, o): return self
    def __mul__(self, o): return self
    def __rmul__(self, o): return self
    def __truediv__(self, o): return self
    def __eq__(self, o): return isinstance(o, _AutoMock)
    def __hash__(self): return hash(self._name)
    def __lt__(self, o): return True
    def __le__(self, o): return True
    def __gt__(self, o): return False
    def __ge__(self, o): return False
    def __mro_entries__(self, bases): return (object,)

y = _AutoMock('y')
nums = [3, -1, 4, -1, 5]

The last line uses :=, the walrus operator (Python 3.8 and later). It saves a value in a variable inside the expression, so you can test it and keep it without computing it twice.

4. Transform: if/else before for

An if/else before the for does not filter. It chooses a value for every item, so every item comes out.

Where the if goesExampleWhat it does
after for[x for x in xs if x > 0]filters: some items are left out
before for, with else[x if x > 0 else 0 for x in xs]transforms: every item comes out
python
[n if n > 0 else 0 for n in nums]                # [3, 0, 4, 0, 0]
["odd" if n % 2 else "even" for n in nums]       # ['odd', 'odd', 'even', 'even', 'even']
[abs(n) for n in nums if n]                      # [3, 1, 4, 2]  (change the items that pass the filter)
+ setup added so this can run · re-runs 1 earlier block on this page first
nums = [3, -1, 4, -1, 5]

# Earlier blocks on this page that this one builds on, re-run quietly first.
_pm_earlier = [
    'nums = [3, -1, 4, 0, -2]\n[n * n for n in nums]              # [9, 1, 16, 0, 4]\n{n % 2 for n in nums}              # {0, 1}\n{n: n * n for n in nums[:2]}       # {3: 9, -1: 1}\nsum(n * n for n in nums)           # 30  (a generator: no list is built)\ntuple(n * 2 for n in nums[:3])     # (6, -2, 8)  (to get a tuple, pass a generator to tuple())',
]
def _pm_run_earlier(_srcs, _g):
    import sys as _s, io as _io, warnings as _w
    for _src in _srcs:
        _o, _e = _s.stdout, _s.stderr
        _s.stdout = _s.stderr = _io.StringIO()
        try:
            with _w.catch_warnings():
                _w.simplefilter('ignore')
                exec(compile(_src, '<earlier block>', 'exec'), _g)
        except (Exception, SystemExit):
            pass
        finally:
            _s.stdout, _s.stderr = _o, _e
    if 'matplotlib.pyplot' in _s.modules:
        _s.modules['matplotlib.pyplot'].close('all')   # their plots belong to their own blocks
_pm_run_earlier(_pm_earlier, globals())
del _pm_earlier, _pm_run_earlier

5. Text

An empty string counts as False, so if w skips empty strings.

python
words = ["Apple", "", "fig", "Kiwi "]
[w.strip().lower() for w in words if w]          # ['apple', 'fig', 'kiwi']
[len(w) for w in words]                          # [5, 0, 3, 5]
[int(x) for x in "3,1,2".split(",")]             # [3, 1, 2]  (turns text into numbers)
[r.split(",") for r in ["a,1", "b,2"]]           # [['a', '1'], ['b', '2']]  (split each line into fields)
[c for c in "a1b2" if c.isdigit()]               # ['1', '2']  (a loop over a string gives one character at a time)
[i for i, w in enumerate(words) if w]            # [0, 2, 3]  (the positions of the words that are not empty)

6. Nested comprehensions

With two for clauses, read them from left to right, in the same order as two nested loops. A comprehension inside another comprehension keeps the rows.

python
grid = [[1, 2, 3], [4, 5, 6]]
[x for row in grid for x in row]                 # [1, 2, 3, 4, 5, 6]  (flatten the grid)
[x for row in grid for x in row if x % 2]        # [1, 3, 5]
[[x * 10 for x in row] for row in grid]          # [[10, 20, 30], [40, 50, 60]]
[list(col) for col in zip(*grid)]                # [[1, 4], [2, 5], [3, 6]]  (rows become columns)
[(a, b) for a in "xy" for b in (1, 2)]           # [('x', 1), ('x', 2), ('y', 1), ('y', 2)]
[[0] * 3 for _ in range(2)]                      # [[0, 0, 0], [0, 0, 0]]

In the last line, _ is a normal variable name that says "I do not use this value".

7. Dict and set comprehensions

python
menu = {"tea": 3, "cake": 5, "pie": 4}
{k: v * 2 for k, v in menu.items()}              # {'tea': 6, 'cake': 10, 'pie': 8}  (change every price)
{k: v for k, v in menu.items() if v > 3}         # {'cake': 5, 'pie': 4}  (keep some pairs)
{v: k for k, v in menu.items()}                  # {3: 'tea', 5: 'cake', 4: 'pie'}  (swap keys and values)
{name: len(name) for name in menu}               # {'tea': 3, 'cake': 4, 'pie': 3}  (a loop over a dict gives its keys)
{k: v for k, v in zip("ab", [1, 2])}             # {'a': 1, 'b': 2}  (dict(zip(...)) does the same)
sorted({len(name) for name in menu})             # [3, 4]

8. Generators passed to a function

When a generator is the only argument of a function, you do not need a second pair of brackets. Functions like sum(), any() and max() take the values one at a time, so no list is built. Not building a list saves memory when there are millions of values.

python
words = ["fig", "kiwi", "lemon"]
max(len(w) for w in words)                       # 5
any(w.startswith("k") for w in words)            # True  (stops at the first True)
all(w.islower() for w in words)                  # True
sum(1 for w in words if len(w) > 3)              # 2  (count the matches)
next(w for w in words if "m" in w)               # 'lemon'  (the first match)
next((w for w in words if "z" in w), None)       # None  (a default when nothing matches)
max((len(w) for w in []), default=0)             # 0

9. Useful patterns

python
[a + b for a, b in zip([1, 2], [10, 20])]        # [11, 22]  (add two lists item by item)
s = "abcde"
[s[i:i + 2] for i in range(0, len(s), 2)]        # ['ab', 'cd', 'e']  (cut the string into pieces of 2 characters)
list(dict.fromkeys([3, 1, 3, 2]))                # [3, 1, 2]  (remove duplicates and keep the order)
[str(n) for n in range(3)]                       # ['0', '1', '2']  (turn every number into a string)
["pos" if x > 0 else "neg" if x < 0 else "zero" for x in [2, -1, 0]]   # ['pos', 'neg', 'zero']

The last line chains two if/else parts. Read it as: "pos" if x > 0, else "neg" if x < 0, else "zero".

Common mistakes

Putting else in the filter. [n for n in nums if n > 0 else 0] is a SyntaxError. When you need else, the if goes before the for: [n if n > 0 else 0 for n in nums].

Expecting round brackets to make a tuple. (n * 2 for n in nums) is a generator. Write tuple(n * 2 for n in nums).

Using a generator twice. A generator is empty after one loop over it:

python
g = (n for n in [3, -1])
list(g)          # [3, -1]
list(g)          # []

Building a grid with *. [[0] * 2] * 2 puts the same row object in the list twice, so a change to one row shows in both:

python
g = [[0] * 2] * 2
g[0][0] = 1
g                # [[1, 0], [1, 0]]
g = [[0] * 2 for _ in range(2)]
g[0][0] = 1
g                # [[1, 0], [0, 0]]

Using the loop variable after the comprehension. The variable exists only inside the brackets. After [m for m in range(3)], the name m is not defined, and using it raises NameError (unless an earlier line created m).

Using a comprehension only for its side effects. [print(n) for n in nums] prints, but also builds a list of None values that nobody uses. Write a normal for loop.

Writing one that is too hard to read. If a comprehension needs three for clauses and two ifs, write it as a normal for loop. The loop is easier to read and to debug.

See also

Practice this

on practicepython.in

Short exercises that run in your browser and tell you what your code did, as well as whether a test passed.

Found this useful? Share it with a friend who’s learning Python.