Skip to content
wannacode

List comprehensions and generators

Problem

A report in a loop

Orders now have a city too:

Order
class Order:
    def __init__(self, customer, city, total):
        self.customer = customer
        self.city = city
        self.total = total
orders = [
    Order("Maria", "Sofia", 75.50),
    Order("Georgi", "Varna", 12.90),
    Order("Elena", "Sofia", 40.00),
    Order("Dimitar", "Plovdiv", 120.00),
    Order("Ana", "Varna", 64.00),
]

You want the customers with orders of at least €50, in alphabetical order:

result = []
for o in orders:
    if o.total >= 50:
        result.append(o.customer)
result.sort()
print(result)
Output
['Ana', 'Dimitar', 'Maria']

Six lines, and the question is just one. The code says how: make an empty list, loop, check, add, sort. To understand what it does, you have to run it in your head. Add "only the first three" or "no duplicates", and the loop will grow.

The idea

A list comprehension: what, not how

The same report, on one line:

result = sorted([o.customer for o in orders if o.total >= 50])
print(result)
Output
['Ana', 'Dimitar', 'Maria']

[o.customer for o in orders if o.total >= 50] is called a list comprehension. It reads almost like a sentence: "the customer of each order o in orders, if its total is at least 50". The parts are:

  • for o in orders: where the elements come from;
  • if o.total >= 50: which ones stay. This is a condition, as in filter. It can also be left out;
  • o.customer at the start: what each element turns into, as in map;
  • the square brackets: the result is a new list.

sorted(...) sorts and also returns a new list. The list orders does not change.

More precisely, if you are curious
  • The same steps also exist as functions with lambdas from the previous lesson: filter and map. The result is the same, but it reads from the inside out, which is why in Python you usually write a list comprehension:
result = sorted(map(lambda o: o.customer, filter(lambda o: o.total >= 50, orders)))
print(result)
Output
['Ana', 'Dimitar', 'Maria']
  • If the expression is the only parameter of a function, the square brackets can be left out: sorted(o.customer for o in orders if o.total >= 50) gives the same. Then the expression is a generator. The next card is about it.

See it

What gets through the filter

Here the steps are two generators: expressions in round brackets instead of square ones. A generator does not make a list. It gives the values one at a time, when something asks for them, and only then computes them.

prices = [12, 75, 40, 130, 58]
expensive = (p for p in prices if p >= 50)
labels = (f"€{p}" for p in expensive)
print(list(labels))
Step by step, in words
  1. Five prices. The two generators only set up the path: nothing has gone through yet.
  2. 12 does not pass if p >= 50 and drops out. It never gets turned into text.
  3. 75 passes if, f"€{p}" turns it into the text "€75", list() collects it.
  4. 40 drops out too. Each element goes the whole way before the next one starts.
  5. 130 and 58 pass. The list prices does not change.
  6. The result is a new list with three strings.

The five prices go one at a time. 12 gets to the check if p >= 50 and drops out. 75 passes, the second generator turns it into the text "€75", and list() collects it. Then 40 drops out, and 130 and 58 pass.

Each element goes the whole way before the next one starts. And nothing starts until list() asks for the values.

More precisely, if you are curious

This is called lazy evaluation. The generator only records what needs to be done. See it with a print in every step:

def check(p):
    print("filter", p)
    return p >= 50




def double(p):
    print("map", p)
    return p * 2




stream = (double(p) for p in [12, 75, 40] if check(p))
print("The generator is ready, nothing has gone through it yet.")
print(list(stream))
Output
The generator is ready, nothing has gone through it yet.
filter 12
filter 75
map 75
filter 40
[150]
  • The line "The generator is ready" is printed before a single element has gone through the generator.
  • map 75 comes right after filter 75: the elements do not wait for all of them to pass the check.
  • With square brackets the same expression is a list and is computed right away, on the line where it is written:
def check(p):
    print("filter", p)
    return p >= 50




def double(p):
    print("map", p)
    return p * 2




now = [double(p) for p in [12, 75, 40] if check(p)]
print("The list is ready.")
print(now)
Output
filter 12
filter 75
map 75
filter 40
The list is ready.
[150]
  • That is why you can take just the first few values from a generator without checking the rest. islice from the itertools module stops as soon as it has taken as many as it needs:
from itertools import islice




def check(p):
    print("filter", p)
    return p >= 50




stream = (p for p in [75, 130, 58, 12, 99] if check(p))
print(list(islice(stream, 2)))
Output
filter 75
filter 130
[75, 130]

The idea

Sum, count, average

The result can be a number instead of a list. The sum of the even numbers:

numbers = [1, 2, 3, 4, 5]
total = sum(n for n in numbers if n % 2 == 0)
print(f"Sum of the even numbers: {total}")
Output
Sum of the even numbers: 6
sum adds up the numbers it gets. The generator inside it has no extra brackets, because it is its only parameter.

The same for the orders:

big = len([o for o in orders if o.total >= 50])
revenue = sum(o.total for o in orders)
average = revenue / len(orders) if orders else 0
max_total = max((o.total for o in orders), default=0)


print(f"Orders over €50: {big}")
print(f"Revenue: €{revenue:.2f}")
print(f"Average: €{average:.2f}")
print(f"Largest: €{max_total:.2f}")
Output
Orders over €50: 3
Revenue: €312.40
Average: €62.48
Largest: €120.00
  • len gives the count. Here we count the elements of a list comprehension.
  • The average is the sum divided by the count. If there are no orders, dividing by 0 stops the program, hence if orders else 0.
  • max of an empty list stops with a ValueError. default=0 says what it should be if there is no value. That is the next idea.
More precisely, if you are curious

With an empty list:

prices = []
print(sum(prices))
print(max(prices, default=0))
try:
    print(max(prices))
except ValueError as e:
    print("ValueError:", e)
try:
    print(sum(prices) / len(prices))
except ZeroDivisionError as e:
    print("ZeroDivisionError:", e)
Output
0
0
ValueError: max() iterable argument is empty
ZeroDivisionError: division by zero

The idea

Sort, cut, no duplicates

ratings = [5, 3, 5, 4, 2, 4, 5]
print(sorted(set(ratings)))


top = [o.customer for o in sorted(orders, key=lambda o: o.total, reverse=True)[:3]]
print(top)


for i in range(1, 4):
    print(f"Table {i}")
Output
[2, 3, 4, 5]
['Dimitar', 'Maria', 'Ana']
Table 1
Table 2
Table 3
  • set(...) removes the duplicates: in a set each value is there once. The order in a set is not guaranteed, so sorted comes after it.
  • sorted(orders, key=..., reverse=True) sorts from the largest order, with key from the previous lesson.
  • [:3] takes the first three elements. This is a slice: from the start up to index 3, without it. Together with sorted, that is a "top 3".
  • range(1, 4) gives the numbers 1, 2, 3. The end is not included.
More precisely, if you are curious

If you need to remove the duplicates but keep the order of first appearance, dict.fromkeys makes a dictionary with the keys in the order they were added:

ratings = [5, 3, 5, 4, 2, 4, 5]
print(list(dict.fromkeys(ratings)))
Output
[5, 3, 4, 2]

Try it

What gets to the end?

What will the program print?

drinks = ["coffee", "tea", "coffee", "water", "tea", "juice"]
result = [d.upper() for d in sorted(set(drinks))[:3]]
print(result)
Answer

The idea

None: a result that may be missing

Which order is the largest? If the list is empty, there is no answer. Then the result is None: "nothing". But the caller has to remember to check. If they forget, the program stops with an AttributeError somewhere later.

Java has Optional for this: a box that either holds a value or is empty. Python has no such box. A missing value is None, and checking it is your job, at the place where the result is used:

biggest = max(orders, key=lambda o: o.total, default=None)
if biggest is not None:
    print(f"Largest: {biggest.customer}")


huge = next((o for o in orders if o.total > 500), None)
print(huge is not None)
print(huge.customer if huge is not None else "no such order")
Output
Largest: Dimitar
False
no such order
  • max(..., default=None) returns None instead of stopping if the list is empty.
  • next(generator, None) gives the first value of the generator, or None if there is none. This is "the first order over €500", without checking the rest.
  • is not None checks whether there is a value. You compare with None using is, not ==: that is what PEP 8 asks for.
  • x if condition else y gives the fallback if there is no value.

Your own functions also return None when they have no answer:

def nickname(user):
    # Returns the nickname, or None if the person has none.
    return "Mimi" if user == "maria" else None
maria = nickname("maria")
if maria is not None:
    print(f"Nickname: {maria}")


georgi = nickname("georgi")
print(georgi if georgi is not None else "no nickname")
Output
Nickname: Mimi
no nickname
More precisely, if you are curious
  • If you forget the check:
huge = next((o for o in orders if o.total > 500), None)
print(huge.customer)
Error
Traceback (most recent call last):
  ...
AttributeError: 'NoneType' object has no attribute 'customer'
  • x or y also gives a fallback, but it treats an empty string and 0 as missing too. Use it only if an empty string really means missing:
print(None or "no nickname")
print("" or "no nickname")
print(0 or 5)
Output
no nickname
no nickname
5
  • any(...) and all(...) answer with True or False whether any or all of the elements match the condition:
print(any(o.total > 100 for o in orders))
print(all(o.total >= 10 for o in orders))
Output
True
True

Try it

Find the bug

The program should print how many prices are expensive, and the prices themselves. The count comes out right, but the prices are an empty list. Which line is the bug on?

Pick the line with the bug.

The idea

Grouping and joining

How many orders are there from each city, and what is the revenue?

from collections import Counter, defaultdict


count_by_city = Counter(o.city for o in orders)
print(dict(sorted(count_by_city.items())))


revenue_by_city = defaultdict(float)
for o in orders:
    revenue_by_city[o.city] += o.total
for city, total in sorted(revenue_by_city.items()):
    print(f"{city}: €{total:.2f}")


customers = ", ".join(sorted(o.customer for o in orders))
print(customers)
Output
{'Plovdiv': 1, 'Sofia': 2, 'Varna': 2}
Plovdiv: €120.00
Sofia: €115.50
Varna: €76.90
Ana, Dimitar, Elena, Georgi, Maria
  • Counter from the collections module counts: for each city, how many times it appears. It is a dictionary.
  • defaultdict(float) starts each new city from 0.0 (that is what float() returns), and += adds the total.
  • ", ".join(...) joins the strings into one, with a comma between them.

The dictionary keeps the keys in the order they were added. sorted(dictionary.items()) gives the key and value pairs sorted by key, so the report comes out alphabetically.

More precisely, if you are curious
  • For each city, a list of the names:
from collections import defaultdict


names_by_city = defaultdict(list)
for o in orders:
    names_by_city[o.city].append(o.customer)
print(dict(sorted(names_by_city.items())))
Output
{'Plovdiv': ['Dimitar'], 'Sofia': ['Maria', 'Elena'], 'Varna': ['Georgi', 'Ana']}
  • A dictionary comprehension makes a dictionary from a list: {key: value for ...}. If a key repeats, the last value wins, with no error (in Java toMap throws an exception):
last_by_city = {o.city: o.customer for o in orders}
print(last_by_city)
Output
{'Sofia': 'Elena', 'Varna': 'Ana', 'Plovdiv': 'Dimitar'}

Exercise

Reports with list comprehensions

Solve the tasks in your editor.

1. Sales by category. A class Sale with a product, a category, a quantity and a unit price, and a method revenue() (quantity times price). Add up the revenue by category with defaultdict(float), print the categories alphabetically and the total revenue. The sales: coffee beans, drinks, 12 at €18.50; tea, drinks, 30 at €4.20; chocolate, sweets, 40 at €2.80; biscuits, sweets, 25 at €3.10; honey, food, 8 at €12; coffee beans again, 5 at €18.50.

Output
Drinks: €440.50
Food: €96.00
Sweets: €189.50
Total: €726.00

2. Top 3 products. With the same sales, find the three products with the highest revenue. The coffee is in two sales, so first group by product and add up, then sort the map's entries by value. Hint: sorted(by_product.items(), key=lambda item: item[1], reverse=True)[:3]: each element of items() is a pair (product, revenue).

Output
Coffee beans: €314.50
Tea: €126.00
Chocolate: €112.00

3. Average grade. Read one line of grades separated by spaces. Print the average, the highest and how many are excellent (5.50 and up). Hint: [float(g) for g in input().split()].

Input
4.50 5.75 6.00 3.00 5.50
Output
Average: 4.95
Highest: 6.00
Excellent: 3

4. The words in a text. Read a line of text and print the words in it with no duplicates, in alphabetical order, in lowercase, separated by commas. Find the runs of letters with re.findall(r"[^\W\d_]+", text): \w is a letter, a digit or _, and [^\W\d_] is "a letter, but not a digit and not _".

Input
Coffee with milk, coffee without sugar. Tea with honey!
Output
coffee, honey, milk, sugar, tea, with, without

More tasks

  • The longest word. With max(words, key=len), find the longest word in a sentence. What will happen for an empty line, and how does default help?
  • Average by class. Students with a name, a class and a grade: collect the grades for each class in a defaultdict(list) and print the average for each class.
  • Big and small orders. With any, check whether there is an order over €1000, and with all whether all of them are at least €10.

Test

Check yourself

1. What happens to the list orders after [o for o in orders if o.total > 50]?
2. What is the difference between [p * 2 for p in prices] and (p * 2 for p in prices)?
3. What does max(orders, key=lambda o: o.total) do with an empty list?
4. Why does the report by city use sorted(revenue_by_city.items())?
5. expensive = (p for p in prices if p >= 50); print(len(list(expensive))); print(list(expensive)). What does the second line print?

Made by Ivo Gergov as part of Evo Tech's campaign for accessible programming education ·