Module 3.1 introduced lists and dictionaries as core data structures. Module 3.3 showed how Pandas and NumPy process data at scale. Between those layers sits a Python idiom you will see in every notebook, API wrapper, and data pipeline: comprehensions—concise expressions that build lists and dicts from iterables.
AI engineers use comprehensions to filter training examples, reshape API responses, build feature dictionaries, and preprocess text before tokenization. They are faster to read than nested loops once you learn the syntax, and they often outperform equivalent imperative code because Python optimizes them internally.
Learning Objectives
By the end of this lesson, students should be able to:
- Write list comprehensions with optional
iffilters and nested loops. - Build dictionary comprehensions for key–value transformations.
- Convert between loop-based and comprehension-based code without changing behavior.
- Choose comprehensions over loops when clarity is preserved, and avoid them when logic is complex.
- Apply comprehensions to common AI tasks: filtering datasets, mapping labels, and reshaping JSON.
- Recognize set comprehensions as a related pattern for unique collections.
Introduction: Loops in One Expression
A list comprehension creates a new list by applying an expression to each item in an iterable, optionally filtering items along the way. Instead of initializing an empty list and appending in a loop, you declare the transformation in a single line.
In machine learning workflows, comprehensions appear constantly: extracting text fields from JSON records, keeping only rows above a confidence threshold, or converting class indices to one-hot positions. Mastering this syntax makes Pandas, PyTorch DataLoader code, and FastAPI response shaping far more readable.
[expression for item in iterable]
With a filter: [expression for item in iterable if condition]
With nested iteration: [expression for a in iter_a for b in iter_b]
Basic List Comprehensions
# Loop equivalent
squares = []
for n in range(10):
squares.append(n ** 2)
# Comprehension
squares = [n ** 2 for n in range(10)]The comprehension version states intent directly: “a list of n squared for each n in range(10).” No temporary variable, no repeated append calls.
After running a classifier, you may keep only high-confidence predictions:
results = [
{"text": "invoice due", "score": 0.92},
{"text": "hello team", "score": 0.31},
{"text": "payment received", "score": 0.88},
]
high_conf = [r["text"] for r in results if r["score"] >= 0.8]
# ["invoice due", "payment received"]Dictionary Comprehensions
Dictionary comprehensions follow the same pattern but produce key–value pairs: {key_expr: value_expr for item in iterable}.
labels = ["cat", "dog", "cat", "bird"]
counts = {label: labels.count(label) for label in set(labels)}
# {"cat": 2, "dog": 1, "bird": 1}
# Map class names to integer IDs for a model
class_to_id = {name: i for i, name in enumerate(sorted(set(labels)))}Dictionary comprehensions are ideal for building lookup tables, inverting mappings, and normalizing API payloads before feeding them to a model.
| Pattern | Syntax | Typical AI Use |
|---|---|---|
| List comprehension | [f(x) for x in xs] | Transform features, token lengths, scores |
| Filtered list | [x for x in xs if cond(x)] | Keep valid rows, non-empty strings |
| Dict comprehension | {k: v for ...} | Label maps, config dicts from env pairs |
| Set comprehension | {x for x in xs} | Unique tokens, unique entity IDs |
When to Use—and When to Avoid
Good Fit
- Simple mapping or filtering over a collection
- One or two levels of nesting
- Logic fits on one readable line
Poor Fit
- Side effects (file writes, API calls) inside the expression
- Deep nesting that obscures intent
- Complex branching better expressed as a function
Reality: They are often slightly faster for simple cases, but readability matters more. For heavy numerical work, use NumPy vectorization. For large dataframes, use Pandas methods like df.apply() or built-in column operations.
Knowledge Check
- Write: Create a list of squares for even numbers 0–9. Answer:
[n**2 for n in range(10) if n % 2 == 0] - True/False: Comprehensions can include
iffilters. Answer: True. - Short Answer: What does
{w: len(w) for w in words}produce? Answer: A dict mapping each word to its length. - Multiple Choice: Best practice for API calls inside iteration: (a) list comprehension, (b) explicit loop, (c) dict comprehension, (d) set comprehension. Answer: (b).
- Write: Invert
{"a": 1, "b": 2}with a comprehension. Answer:{v: k for k, v in d.items()} - Short Answer: Write the basic list-comprehension pattern. Answer:
[expression for item in iterable]. - True/False: Set comprehensions use curly braces to collect unique values. Answer: True (
{x for x in xs}). - Multiple Choice: Nested comprehensions with deep branching are: (a) always best, (b) often a poor fit—extract a function, (c) required by NumPy, (d) faster than GPUs. Answer: (b).
- Short Answer: How do you keep only high-confidence prediction texts? Answer:
[r["text"] for r in results if r["score"] >= 0.8](or equivalent). - True/False: Comprehensions are always faster than loops, so use them even for API calls. Answer: False—readability first; avoid side effects inside comprehensions.
Key Takeaways
- List comprehensions build lists in a single expressive statement.
- Dictionary comprehensions create mappings from iterables.
- Use filters (
if) to select subsets during construction. - Prefer clarity over cleverness—extract complex logic into functions.
- Next: Lambda Functions for small anonymous functions often used with
sorted()andmap().
Live demo: Show the same task as a loop, then refactor to a comprehension. Ask students which version states intent faster.
Exercise: Given a list of chat messages with role and content, use a comprehension to extract only user messages longer than 20 characters.
Recap: List and dict comprehensions express mapping and filtering in one line; next, compact callbacks in Lambda Functions.