Python 函式抽取前先用 AST 盤點外部依賴

Most extractions do not fail in review. They fail two weeks later, when a test hits a module global nobody remembered. The repair costs more than the extraction did.

So start with a count, not a diff. Census every name the target function reads. Then move exactly that much and no more.

Why the import list lies

An import header tells you what a file names. It does not tell you what a function reads at runtime. Module globals, re-exported symbols, and monkeypatched hooks never show up there.

The census reads the function body instead. It walks the AST and collects every name the body could read from outside itself. That set is the real dependency surface you have to carry.

Step 1 — Pin behavior before you measure

Behavior first. A census of an untested function maps a moving target. Keep this step small and boring.

1. Commit one real input from the repo into the test tree.
2. Freeze the clock, the random seed, and the output sort order.
3. Call the entry point in a subprocess, so import-time side effects are visible.
4. Assert on serialized bytes, not on live objects.

The pass condition matters more than coverage. The test must be able to fail. If a one-byte change in output still passes, the test is decoration. Fix that before you measure anything.

Step 2 — Build the name-touch census

Drop this file at tools/census.py. It answers one question per function: which module-level names does this body read?

# tools/census.py
import ast, builtins, json, sys
from pathlib import Path

BUILTINS = set(dir(builtins))


def free_names(fn):
    """Names a function body reads that are not local to it."""
    local = {a.arg for a in fn.args.posonlyargs + fn.args.args + fn.args.kwonlyargs}
    for extra in (fn.args.vararg, fn.args.kwarg):
        if extra:
            local.add(extra.arg)
    reads, writes = set(), set()
    for node in ast.walk(fn):
        if isinstance(node, ast.Name):
            (writes if isinstance(node.ctx, ast.Store) else reads).add(node.id)
        elif isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)) and node is not fn:
            writes.add(node.name)
        elif isinstance(node, (ast.Import, ast.ImportFrom)):
            for a in node.names:
                writes.add((a.asname or a.name).split(".")[0])
        elif isinstance(node, ast.comprehension):
            for n in ast.walk(node.target):
                if isinstance(n, ast.Name):
                    writes.add(n.id)
    return reads - local - writes


def module_bindings(tree):
    """What the module itself binds at top level."""
    bound = {}
    for node in tree.body:
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
            bound[node.name] = "def"
        elif isinstance(node, ast.Assign):
            for target in node.targets:
                for n in ast.walk(target):
                    if isinstance(n, ast.Name):
                        bound[n.id] = "assign"
        elif isinstance(node, (ast.Import, ast.ImportFrom)):
            for a in node.names:
                bound[(a.asname or a.name).split(".")[0]] = "import"
    return bound


def main(path, func_name):
    tree = ast.parse(Path(path).read_text())
    bound = module_bindings(tree)
    fn = next(n for n in ast.walk(tree)
              if isinstance(n, (ast.FunctionDef, ast.AsyncFunctionDef))
              and n.name == func_name)
    free = free_names(fn)
    print(json.dumps({
        "file": path,
        "function": func_name,
        "allowed_imports": sorted(n for n in free if bound.get(n) == "import"),
        "module_globals": sorted(n for n in free if n in bound and bound[n] != "import"),
        "unbound": sorted(n for n in free if n not in bound and n not in BUILTINS),
        "builtins_read": sorted(n for n in free if n in BUILTINS),
    }, indent=2))


if __name__ == "__main__":
    main(sys.argv[1], sys.argv[2])

Run it against the function you intend to extract:

python tools/census.py pkg/report.py pin_report > census.json

Example output, from a report writer that reads a template and a formatter:

{
  "file": "pkg/report.py",
  "function": "pin_report",
  "allowed_imports": ["json", "os"],
  "module_globals": ["DEFAULT_FORMAT", "_formatter"],
  "unbound": ["SINK"],
  "builtins_read": ["len", "sorted"]
}

Read unbound first. That list is where extractions actually break. SINK is not bound anywhere in the module, so something injects it at runtime. Find the writer before you move the function.

Step 3 — Turn the census into a budget

The counts decide the shape of the change. Four buckets, four actions.

The third row carries the real lesson. A function that reads a mutable module global is not extractable yet. Make the read explicit at the call site, run the pinned test, then move.

Step 4 — Move, then shim

The smallest safe change is a move plus a re-export. No call site changes. The diff touches two files and zero callers.

# pkg/report_writer.py — new home, imports limited to the census list
import json
import os


def pin_report(rows, fmt):
    ...

# pkg/report.py — old home keeps the name alive
from pkg.report_writer import pin_report  # shim; delete after callers migrate

Verify the pinned test passes twice: once against the shim, once against the new module. Same bytes, both paths. Then migrate callers in a separate change, not this one.

Step 5 — Gate the import surface in CI

The new module must not import more than the census allowed. Enforce it mechanically, not by memory.

# tools/gate.py — fails when the extracted module imports beyond the census
import ast, json, sys

census = json.load(open("census.json"))
allowed = set(census["allowed_imports"])
tree = ast.parse(open(sys.argv[1]).read())
imports = set()
for node in ast.walk(tree):
    if isinstance(node, ast.Import):
        imports |= {(a.asname or a.name.split(".")[0]) for a in node.names}
    elif isinstance(node, ast.ImportFrom):
        imports |= {(a.asname or a.name) for a in node.names}
added = imports - allowed
print("allowed:", sorted(allowed))
print("added:  ", sorted(added))
sys.exit(1 if added else 0)

Wire it into the pipeline:

python tools/census.py pkg/report.py pin_report > census.json
python tools/gate.py pkg/report_writer.py

When someone later slips requests into the extracted module, the gate exits 1 and names the addition. That is a cheaper failure than a staging incident.

Where a free model fits, and where it must not

Disclosure: This article was prepared as part of MonkeyCode's product outreach.

The census output is a good short prompt. Feed the JSON and the entry point signature to a free model and ask for candidate characterization tests. Models are decent at enumerating edge cases you skip: empty input, unicode separators, zero rows.

MonkeyCode's free model access covers that drafting step without a spend approval. Treat the drafts as candidates. Every suggested test still has to fail on a mutated build before it joins the suite.

A second environment helps too. If your local Python minor version differs from CI's, run the census and the gate on MonkeyCode's free server option against a clean checkout. That catches import-time assumptions your laptop silently satisfies.

Do not ask a model to decide the extraction boundary. It cannot see your runtime globals, and that is exactly what the census exists to find.

Limitations and who should skip this

The census is a heuristic. It will mislead you in four known ways.

• Dynamic reads through globals(), getattr, or eval stay invisible.
• Inner functions are approximated, so a nested scope can hide a read.
• Decorators that rewrite signatures confuse the local-variable set.
• Star imports collapse many names into one binding.

Skip the whole workflow when the function is already pure and under twenty lines. Just move it. Skip it too when the repo has no runnable entry point: building one is the real task, and the census will not help.

The rule is narrow and it holds: count the names, spend the budget, move once. If you want to try the drafting step first, the free model access is the cheapest place to start.