Getting the most out of Python collections
Tags: python, collections, tips
Published: 2020-02-18 09:17:33

---
title: Getting the most out of Python collections
description: A guide to comprehensions, generators and useful functions and classes
published: true
date: 2020-02-18 00:00:00 UTC
tags: python, collections, tips
canonical_url: https://sourcery.ai/blog/effective-collection-handling/?utm_souce=Blog&utm_medium=RSS
cover_image: https://dev-to-uploads.s3.amazonaws.com/i/8hmhi776rr8uxsfqfo75.jpeg
---

One of Python's best features is that it has awesome capabilities for creating and handling collections. Building a good understanding of these will help you write clean Pythonic code.

To get you started this guide covers comprehensions, generators, and some really useful built-in functions and collection types.

# Create and filter collections with Comprehensions

List, Dictionary and Set comprehensions are a powerful feature of the Python language that enable you to declare these collection types and fill them with values at the same time. This lets you do all sorts of filtering and mapping operations in a simple and readable way.

## List Comprehensions

Without comprehensions a common pattern used to create a list and fill it with values is as follows:

```python
cubes = []
for i in range(20):
  cubes.append(i ** 3)
```

Each operation is carried out separately - an assignment to create the empty list, and then a for loop iterating over a sequence within which the list is appended to. We end up with a list of cubed numbers sitting in the `cubes` variable.

This can be shortened to the following using a list comprehension:

```python
cubes = [i ** 3 for i in range(20)]
```

Let's unpack what's going on. At first glance it can be confusing, with lots of the elements from the for loop version moved into different places in the code.

We know the purpose of the code is to create a list of cubed numbers. We therefore have an assignment using square brackets to signify that we are creating a list. Inside the comprehension we can start with the second half: `for i in range(20)`. This is a standard for loop, but instead of a for body we put an expression in front of it giving the element that we will create at each iteration. In this case we have `i ** 3`, cubing the target variable `i`. The expression can be simple, or can involve any sorts of function calls or operations that you like. It does not even have to use the target variable.

We can also add an if clause at the end to filter the list:

```python
odd_cubes = [i ** 3 for i in range(20) if i % 2 != 0]
```

Often squeezing code onto one line can make it more difficult to read. In the case of comprehensions all the elements that you need are nicely presented, and once you are used to the syntax it is actually more readable than the for loop version of the same code.

## Dictionary Comprehensions

You can also directly create dictionaries using a dictionary comprehension. Here you set the keys and values at once using the following syntax:

```python
squares = {x: x ** 2 for x in range(100)}
```

This creates a dictionary mapping numbers to their squared values. As you can see you set the keys and values while iterating over a sequence. Again you can use if clauses in your dictionary comprehensions in order to perform filtering at the same time.

Filtering would look like this:

```python
even_squares = {x: x ** 2 for x in range(100) if x % 2 == 0}
```

## Set Comprehensions

Moving on, there are also set comprehensions, which create a set, so all values are unique.

```python
square_roots = {math.sqrt(number) for number in number_list}
```

Again these can be filtered. If you want to filter using an if expression then the position of the conditional in the comprehension changes like this:

```python
positive_square_roots = {math.sqrt(number) if number > 0 else 0 for number in number_list}
```

## More complex comprehensions

Comprehensions can also replace nested for loops. Take the example of flattening a list of lists:

```python
list_of_lists = [[1, 2], [3, 4]]
flat = []
for sub_list in list_of_lists:
  for element in sub_list:
    flat.append(element)
```

This produces `[1, 2, 3, 4]`. As a comprehension this turns into:

```python
list_of_lists = [[1, 2], [3, 4]]
flat = [element for sub_list in list_of_lists for element in sub_list]
```

You can see that the ordering of the for loops has been maintained. Using comprehensions for nested for loops in this way becomes a bit unintuitive. Past a certain level of complexity it is more difficult to understand than the same code would be using a for loop.

Just because you can achieve something in a one line comprehension using multiple for and if clauses doesn't mean that it's a good idea.

# Iterating with Generators

A generator is a special type of Python function that returns a lazy iterator. You can loop over these results without the whole list of them having to be stored in memory.

For a concrete example take this function that generates the Fibonacci sequence of numbers:

```python
def fib():
    a, b = 0, 1
    while 1:
       yield b
       a, b = b, a + b
```

You might then call it with code like this:

```python
def print_fibonacci(max_size)
  for number in fib():
    if number < max_size:
      print(number)
    else:
      break
```

So what's going on here? From the point of view of the calling function `fib()` just looks like a collection.
When `fib()` is first invoked it executes as normal until it hits the `yield b` statement. It then returns control to its caller
along with the value, just like `return b` would do. However, `fib()` has only been suspended, and all its state has been saved. When the iterator it returns is queried again it resumes from where it left off instead of starting again from the beginning like a normal function.

The other way of defining generators in Python is using generator expressions. These use the same syntax as comprehensions, but are wrapped in round brackets:

```python
cubes = (i ** 3 for i in range(20))
```

These return an iterator rather than a collection. If the iterator is exhausted and you ask it for another item it will return a `StopIteration` exception.

In fact a comprehension behaves like wrapping a generator expression in a call to `list()` or `set()`:

```python
cubes = list(i ** 3 for i in range(20))
```

Note that when passing a generator into a function, Python allows you to omit one of the pairs of brackets (like in the example just above - we are making a call to `list()`).

One tip is that functions like `any`, `all` and `sum` allow you to pass in a generator rather than a collection. This may lead to performance improvements due to the lazy evaluation (i.e. the function may be able to return before evaluating all of the elements from the generator).

The primary benefit of generators comes from the aforementioned lazy evaluation. This can make the memory overhead of a generator much lower than creating the whole collection at once and then iterating over it.

Take the example of reading lines from a file. If the file is large then reading the whole of it into memory may be extremely slow or impossible. The file object returned by `open()` is already a generator of lines, so can be used like this:

```python
def read_lines(file):
  with open(file) as f:
    for line in f:
      print(line)
```

# Useful types of collection from the [`collections`](https://docs.python.org/3.8/library/collections.html) module

An important way to improve your collections code is to make use of the available builtin types. Here are two of the handiest:

## `defaultdict`

This is a subclass of dictionary that calls a factory function to supply missing values.

This lets you replace this:

```python
fruit_sizes = {}
for fruit in fridge:
  if fruit.name in fruit_sizes:
    fruit_sizes[fruit.name].append(fruit.size)
  else:
    fruit_sizes[fruit.name] = [fruit.size]
```

with this:

```python
fruit_sizes = defaultdict(list)
for fruit in fridge:
  fruit_sizes[fruit.name].append(fruit.size)
```

Here any missing values are initialized to empty lists. This avoids cluttering up the code with checks and initializations.

Other common arguments are to create `defaultdicts` of `int`, `float`, `dict` or `set` types.

To create one which sets the default to anything you like you can use lambda functions. For example if all students are assumed to have a small green hat you would use:

```python
student_hats = defaultdict([Hat('green', 'small')])
```

## `Counter`

Need to count occurrences of multiple different hashable objects? A `Counter` is your friend. It is a dictionary subclass where the elements are the keys and their counts are the values.

When creating or updating the counter you give it an iterable. It iterates through and counts the occurrences of items, matching them by their hashes.

An easy way to get started on understanding `Counter` is to count all of the occurrences of strings in a list.

```python
counts = Counter(['Fred', 'Samantha', 'Jean-Claude', 'Samantha'])
```

`counts` then looks like this:

```python
Counter({'Samantha': 2, 'Fred': 1, 'Jean-Claude': 1})
```

You can then update the `Counter` with further iterables, or other instances of `Counter`.

# Useful functions for collection handling

Once you have a collection or a generator the next stage is to do something useful with it. Python contains many useful functions for collection handling that can improve how you write your code.

## Using `any()` and `all()`

A common pattern to establish something about a collection is the following:

```python
any_green = False
for hat in hats:
  if hat.colour == 'green':
    any_green = True
    break
```

This can be replaced with a call to `any()` using a generator:

```python
any_green = any(hat.colour == 'green' for hat in hats)
```

This returns whether any item in the generator satisfies the condition, and will shortcut the evaluation just as the above for loop does with the `break` statement.

`all()` behaves similarly, but returns whether all items satisfy the condition.

## Iterating over collections together with `zip()`

You often find yourself with multiple collections which you need to iterate over, collecting data together from each one.

```python
for name, email in zip(names, emails):
  create_person(name, email)
```

This steps through both collections at the same time, returning a tuple consisting of elements from each one.
Here we have unpacked the tuple into separate values. `zip()` can take as many collections as you like, but will stop when the
shortest one is exhausted. If you'd like to deal with all values you can use `itertools.zip_longest(*iterables, fillvalue=None)`. This will carry on until all collections are exhausted, filling in missing values with the `fillvalue`.

One neat use of the `zip()` functions is to iterate over pairs of elements in the same collection. For example:

```python
differences = [next_elt - elt for elt, next_elt in zip(items, items[1:])]
```

## Chaining over multiple collections

If you need to iterate over multiple collections one at a time, you can use `itertools.chain()`.

```python
for name in itertools.chain(first_name_list, second_name_list):
  create_person(name)
```

This iterates through the first collection until it is exhausted, then moves on to the next and so on.

There are lots more interesting and advanced functions in the `itertools` module, and we will return to them in
a further blog post.