Lambdas, streams and comprehensions

This topic covers the Python tools for working with functions and sequences of values: comprehensions, lambdas, generators, decorators and context managers. "Explain decorators, generators and context managers" is a common question for testers, because pytest (the usual Python test tool) is built on all three. Fixtures (functions that prepare things for a test) and markers (labels that group tests) are decorators, yield fixtures are generators, and pytest.raises is a context manager. Each term is explained below.

Step 1: comprehensions

A comprehension builds a new collection from an existing one in one expression, with an optional filter:

• a list: [e.text for e in elements if e.text]
• a dict: {u["id"]: u for u in users}
• a set: {u["role"] for u in users}
• a generator expression, in parentheses: sum(item["price"] for item in cart), which builds no list at all; it hands out one value at a time.

Read a comprehension as a loop: "for each item, if the condition holds, add the expression to the result". Use comprehensions to transform and filter. Use a plain loop when the body clicks, logs, asserts or needs try/except.

Step 2: lambdas, map, filter, any and all

A lambda is a small function with no name and a single expression, as in key=lambda u: u["age"]. It cannot hold statements such as return or an assignment. map (applies a function to every item) and filter (keeps the items that pass a test) return lazy iterators, which produce values only when asked and can be read only once. reduce (combines all items into one value) lives in functools. Most Python code prefers comprehensions and built-ins such as sum, max, any and all; any and all stop as soon as they know the answer.

from functools import reduce

prices = [120.0, 35.5, 4.0, 1.5]
print([p for p in prices if p < 50])
print({name: len(name) for name in ["Asha", "Ravi"]})
print(sum(p for p in prices), max(prices), any(p > 100 for p in prices))
print(list(map(lambda p: p * 2, prices[:2])), reduce(lambda a, b: a + b, prices))

doubled = map(lambda p: p * 2, prices)
print(list(doubled), list(doubled))      # iterators are used up after one pass

users = [{"name": "Ravi", "age": 31}, {"name": "Asha", "age": 28}]
print([u["name"] for u in sorted(users, key=lambda u: u["age"])])

[35.5, 4.0, 1.5]
{'Asha': 4, 'Ravi': 4}
161.0 120.0 True
[240.0, 71.0] 161.0
[240.0, 71.0, 8.0, 3.0] []
['Asha', 'Ravi']

Step 3: generators and decorators

A function that contains yield is a generator function. Calling it runs nothing yet; it returns a generator. Each next() runs the body up to the next yield, hands out that value and pauses with its local variables kept. yield gives a value and pauses; return ends the function. Generators make one value at a time, so you can scan a huge log file without loading it into memory. Like map, a generator can be read only once.

A decorator is a function that takes a function and returns a replacement, usually a wrapper (a function that calls the original and adds something around it) that adds behaviour before and after the call. @retry(3) above a def is short for func = retry(3)(func). Put functools.wraps on the wrapper so the original function's name survives in logs and reports.

import functools

def non_empty_lines(lines):
    for line in lines:
        if line.strip():
            yield line.strip()

gen = non_empty_lines(["asha\n", "\n", "ravi\n"])
print(next(gen), list(gen), list(gen))

def retry(times):
    def decorate(func):
        @functools.wraps(func)
        def wrapper(*args, **kwargs):
            for attempt in range(1, times + 1):
                try:
                    return func(*args, **kwargs)
                except AssertionError:
                    print(f"attempt {attempt} failed")
                    if attempt == times:
                        raise
        return wrapper
    return decorate

calls = {"n": 0}

@retry(3)
def check_banner():
    calls["n"] += 1
    assert calls["n"] >= 2, "banner not shown yet"
    return "banner shown"

print(check_banner(), check_banner.__name__)

asha ['ravi'] []
attempt 1 failed
banner shown check_banner

A retry decorator is a good interview example, but in a real suite retries can hide flaky tests (tests that pass and fail at random), so use them only for known slow outside systems.

Step 4: context managers

A context manager sets something up and always cleans up, even when the block raises an error. A class does this with __enter__ (runs when the with block starts) and __exit__ (runs when it ends). contextlib.contextmanager turns a generator into one: the code before yield is the set-up and the code after it, in a finally, is the clean-up.

from contextlib import contextmanager

@contextmanager
def in_frame(name):
    print(f"switch to {name}")
    try:
        yield name
    finally:
        print("switch back to main page")

try:
    with in_frame("payment") as frame:
        print("filling card form in", frame)
        raise RuntimeError("card field missing")
except RuntimeError as e:
    print("error:", e)

switch to payment
filling card form in payment
switch back to main page
error: card field missing

pytest fixtures have the same shape: code before yield starts a driver, the yielded value goes to the test, and code after yield quits the driver whether the test passed or failed.

How to answer in an interview

Define each tool in one sentence, then point to pytest: "A decorator wraps a function to add behaviour, like @pytest.fixture. A generator uses yield to give values one at a time, like a yield fixture. A context manager guarantees clean-up with with, like pytest.raises." If asked to write one, a retry or timing decorator with functools.wraps is a safe choice.

Common mistakes

• Reading a generator or map object twice and getting nothing the second time.
• Forgetting functools.wraps, so every decorated test shows the name "wrapper".
• Writing return inside a lambda.
• Using a comprehension only to call a function for its side effect.
• Putting a list comprehension inside any(), which checks every item instead of stopping early.

Recap

Comprehensions transform and filter in one line, and generator expressions do it without building a list. Lambdas are one-expression functions. yield makes a lazy, one-pass generator. Decorators wrap functions, and context managers guarantee clean-up. All of them appear inside pytest.