python-comprehension-refactorer

Turn accumulate-in-a-loop code into comprehensions and generator expressions when that makes it clearer and faster, and turn comprehensions back into loops (or named generator functions) when they have grown past what a reader can hold in their head. The skill covers the whole family: list/set/dict comprehensions, generator expressions, itertools, scoping and leakage rules, side-effect anti-patterns and performance. Its examples come from AI-engineering code: prompt construction, embedding batches, retrieval filtering, tool-call flattening, token counting and dataset streaming.

#Trigger Conditions

Activate this skill when the user:

• Asks to make a loop "more Pythonic", "a one-liner", or "a list/dict comprehension"
• Pastes a comprehension and asks what it does, or says it is unreadable
• Builds lists/dicts in loops in RAG, agent, eval or training code and wants it cleaner
• Hits memory problems from materialising large lists (sum([...]), reading whole files, embedding all texts at once)
• Needs batching/chunking (batched, islice), flattening (chain.from_iterable), grouping (groupby, defaultdict) or sliding windows (pairwise)
• Asks whether map/filter, a comprehension or numpy/pandas vectorization is faster
• Wants a review pass over .py files for comprehension anti-patterns

#Workflow

1. SCAN       -> run scripts/find_comprehension_candidates.py on the file(s)
2. CLASSIFY   -> map, filter, flatten, group, count, reduce, or side effect?
3. CHOOSE     -> list/set/dict comp, generator expr, itertools, Counter/defaultdict, or keep the loop
4. REWRITE    -> apply the suggested code; keep names meaningful; wrap long comps one clause per line
5. CHECK      -> leaked loop variables, class-body scope, exceptions, evaluation order, laziness
6. VERIFY     -> run tests or compare old vs new output on sample data; re-run the script

#Decision Table

#Core Rules

1. A comprehension builds a value. A loop performs actions. Never use a comprehension for side effects ([print(x) for x in xs], [db.add(r) for r in rows]): it allocates a throwaway list of None and hides intent.
2. Two for clauses and two if clauses are the readability limit. Past that, or past one level of nesting (a comprehension inside a comprehension inside a comprehension), unroll into a loop or a named generator function.
3. Clause order equals loop order. [x for row in matrix for x in row] is for row in matrix: for x in row:. The element goes first, everything else reads top-down.
4. Filter if goes at the end, choice if/else goes in front. [x for x in xs if ok(x)] drops items; [a if ok(x) else b for x in xs] keeps every item.
5. Use a generator expression when the result is consumed once. sum(x for x in xs) needs no extra parentheses when it is the sole argument; any/all then short-circuit.
6. Do not compute the same expression twice. [f(x) for x in xs if f(x)] calls f twice; use [y for x in xs if (y := f(x))] (3.8+) or a nested generator [y for y in map(f, xs) if y].
7. Comprehension variables do not leak (Python 3). If code after the loop reads the loop variable, the rewrite changes behaviour. Walrus targets do bind in the enclosing scope.
8. Class bodies are special. Inside a class body only the outermost iterable can see class-level names; the walrus operator in a comprehension is a SyntaxError there.
9. Grouping is not a comprehension. {k: [v for v in vs if key(v) == k] for k in keys} is O(n*k). Use defaultdict(list) in one pass, or itertools.groupby on already-sorted data.
10. Readability beats cleverness. If a comprehension needs a comment to explain it, give it a name instead: a helper function, a generator function, or an intermediate variable.

#Before / After

Prompt list with filtering

# before
prompts = []
for q in questions:
    if q.strip():
        prompts.append(TEMPLATE.format(question=q.strip()))

# after
prompts = [TEMPLATE.format(question=q.strip()) for q in questions if q.strip()]

Retrieved chunks above a score threshold, keeping order

relevant = [hit.text for hit in hits if hit.score >= min_score]
top_k = sorted(hits, key=lambda h: h.score, reverse=True)[:k]   # not a comprehension job

id -> embedding dict, and inversion

vectors = {rec["id"]: rec["embedding"] for rec in records}
label_to_id = {label: idx for idx, label in enumerate(labels)}
id_to_label = {idx: label for label, idx in label_to_id.items()}   # values must be unique

Flatten tool calls from several model responses

calls = [call for resp in responses for call in (resp.tool_calls or [])]
# or, lazily
from itertools import chain
calls = chain.from_iterable(resp.tool_calls or [] for resp in responses)

Parse-and-keep with the walrus operator (3.8+)

answers = [parsed for raw in outputs if (parsed := try_parse_json(raw)) is not None]

Token counting

from collections import Counter
counts = Counter(tok for doc in docs for tok in tokenizer.tokenize(doc))
total_tokens = sum(len(tokenizer.encode(t)) for t in texts)      # generator: no temp list

Batching for an embeddings API

from itertools import batched          # 3.12+; see references for a 3.9 fallback
for batch in batched(texts, 64):
    vectors.extend(client.embed(list(batch)))

Too complex: unroll it

# before: four for clauses and three ifs on one line
toks = [t.lower() for c in corpora for d in c for s in d.sents for t in s if t.is_alpha if not t.is_stop if len(t) > 2]

# after: a named generator function the reader can follow and test
def content_tokens(corpora):
    for corpus in corpora:
        for doc in corpus:
            for sent in doc.sents:
                for tok in sent:
                    if tok.is_alpha and not tok.is_stop and len(tok) > 2:
                        yield tok.lower()

toks = list(content_tokens(corpora))

Side effects belong in a loop

# before
[logger.info("upserting %s", c.id) for c in chunks]
# after
for c in chunks:
    logger.info("upserting %s", c.id)

#Generators and itertools in One Table

Full coverage, including recipes and pitfalls, is in generators-and-itertools.md.

#Scoping and Laziness Gotchas

# leaked variable: before
for chunk in chunks:
    ids.append(chunk.id)
log.info("last chunk %s", chunk.id)

# after: the dependency is explicit
ids = [c.id for c in chunks]
if chunks:
    log.info("last chunk %s", chunks[-1].id)

Details and the full list are in when-not-to-use.md.

#Correctness Checklist

- [ ] Same output order and same duplicates (a set/dict comprehension dedupes)
- [ ] Dict comprehension: duplicate keys keep the LAST value, same as the loop
- [ ] No loop variable read after the loop (comprehension vars do not leak)
- [ ] No break/continue/return/try/except/yield needed inside the body
- [ ] Expensive calls not duplicated between element and condition
- [ ] Generators consumed only once (a second pass yields nothing)
- [ ] Laziness does not outlive a resource (file closed, session ended) before iteration
- [ ] Late binding: lambdas created in a comprehension capture the variable, not the value
- [ ] Not inside a class body when the comprehension reads class attributes or uses walrus

#Performance Facts

• A list comprehension is faster than the equivalent append loop (measured about 8 percent on 3.14, more on 3.9/3.10): it uses a dedicated LIST_APPEND bytecode instead of a method lookup and call per item. PEP 709 (3.12) also inlines list/dict/set comprehensions, removing the hidden function call.
• A generator expression is slower per item than a list for cheap elements (sum(x for x in xs) took 1.5x sum([x for x in xs])); choose it for memory and short-circuiting, not raw speed.
• map(str.strip, lines) with a builtin or C function is as fast as or faster than a comprehension; map(lambda ...) is slower than a comprehension.
• "".join([...]) with a list is slightly faster than with a generator (join materialises it anyway); this is the one reducer where a list is fine.
• For numeric arrays, numpy/torch vectorization is 10 to 100 times faster than any comprehension: scores[scores > t], embeddings @ query, np.linalg.norm(vecs, axis=1).
• Generators trade a little per-item speed for O(1) memory. Use them for streams, files and large corpora, not for 10-element lists you index into.

Numbers, benchmarks and the timing method are in performance-notes.md.

#Output Format

When refactoring, respond with:

1. Findings: each loop or comprehension, its line range and the category (map, filter, flatten, group, count, reduce, side effect, too complex).
2. Rewrite: before and after code for each change, with the target Python version noted when it matters (batched 3.12+, pairwise 3.10+, walrus 3.8+).
3. Behaviour notes: anything that changes (leaked variables, laziness, evaluation count, empty-input errors, ordering, duplicate keys).
4. Left as a loop: what you deliberately did not convert, and why.
5. Verification: how to confirm equivalence (a test, an assert on sample data, or re-running the script).

#Scripts

Find candidates and get exact replacement code:

python scripts/find_comprehension_candidates.py rag/pipeline.py
python scripts/find_comprehension_candidates.py src/**/*.py --json          # machine-readable findings
cat agent.py | python scripts/find_comprehension_candidates.py - --strict    # exit 1 if anything is found (CI)
python scripts/find_comprehension_candidates.py ingest.py --max-for 3 --max-ifs 3 --max-line 120

scripts/find_comprehension_candidates.py parses the code with ast (never executes it) and reports:

Each finding includes line numbers, the suggested code (via ast.unparse), behaviour notes (leaked loop variables, class-body scope, NaN-safe negation) and the reference file to read. Exit codes: 0 ran, 1 findings with --strict, 2 usage, unreadable file or syntax error.

#Additional Resources