Find the Seam Before You Cut: A Read/Write Census for Messy Modules

# seam_census.py -- rank functions by how hard they are to extract safely.
# Python 3.10+, standard library only.
from __future__ import annotations

import ast
import sys
from dataclasses import dataclass, field
from pathlib import Path

MUTATORS = {
    "append", "extend", "insert", "remove", "pop", "clear", "sort",
    "update", "setdefault", "add", "discard", "write", "close",
}


def target_names(node):
    if isinstance(node, ast.Name):
        return {node.id}
    if isinstance(node, (ast.Tuple, ast.List)):
        out = set()
        for elt in node.elts:
            out |= target_names(elt)
        return out
    return set()


def bound_names(stmts):
    """Names bound at module level, plus one level of nesting."""
    out = set()
    for node in stmts:
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef)):
            out.add(node.name)
        elif isinstance(node, ast.Assign):
            for t in node.targets:
                out |= target_names(t)
        elif isinstance(node, (ast.AnnAssign, ast.AugAssign)) and node.target is not None:
            out |= target_names(node.target)
        elif isinstance(node, ast.Import):
            out |= {a.asname or a.name.split(".")[0] for a in node.names}
        elif isinstance(node, ast.ImportFrom):
            out |= {a.asname or a.name for a in node.names}
        elif isinstance(node, (ast.For, ast.AsyncFor)):
            out |= target_names(node.target) | bound_names(node.body) | bound_names(node.orelse)
        elif isinstance(node, (ast.With, ast.AsyncWith)):
            for item in node.items:
                if item.optional_vars is not None:
                    out |= target_names(item.optional_vars)
            out |= bound_names(node.body)
        elif isinstance(node, ast.Try):
            out |= bound_names(node.body) | bound_names(node.orelse) | bound_names(node.finalbody)
            for handler in node.handlers:
                if handler.name:
                    out.add(handler.name)
                out |= bound_names(handler.body)
        elif isinstance(node, (ast.If, ast.While)):
            out |= bound_names(node.body) | bound_names(node.orelse)
    return out


@dataclass
class Seam:
    qualname: str
    lineno: int
    writes: set[str] = field(default_factory=set)
    aliases: set[str] = field(default_factory=set)
    reads: set[str] = field(default_factory=set)
    calls: set[str] = field(default_factory=set)

    @property
    def score(self) -> int:
        return 3 * len(self.writes) + 2 * len(self.aliases) + len(self.reads) + len(self.calls)


def scan(fn, module_names):
    declared = {a.arg for a in fn.args.posonlyargs + fn.args.args + fn.args.kwonlyargs}
    if fn.args.vararg:
        declared.add(fn.args.vararg.arg)
    if fn.args.kwarg:
        declared.add(fn.args.kwarg.arg)
    declared |= bound_names(fn.body)

    seam = Seam(fn.name, fn.lineno)
    declared_global = set()

    for node in ast.walk(fn):
        if isinstance(node, ast.Global):
            declared_global |= set(node.names)
        elif isinstance(node, ast.Name) and node.id in module_names and node.id not in declared:
            if isinstance(node.ctx, ast.Store) and node.id in declared_global:
                seam.writes.add(node.id)
            elif isinstance(node.ctx, ast.Load):
                seam.reads.add(node.id)
        elif isinstance(node, ast.Call):
            func = node.func
            if isinstance(func, ast.Name) and func.id in module_names and func.id not in declared:
                seam.calls.add(func.id)
            elif isinstance(func, ast.Attribute) and func.attr in MUTATORS:
                base = func.value
                if isinstance(base, ast.Name) and base.id in module_names and base.id not in declared:
                    seam.aliases.add(f"{base.id}.{func.attr}()")
        elif isinstance(node, ast.Subscript) and isinstance(node.ctx, ast.Store):
            base = node.value
            if isinstance(base, ast.Name) and base.id in module_names and base.id not in declared:
                seam.aliases.add(f"{base.id}[...] =")

    seam.reads -= seam.writes
    return seam


def main(path: str) -> None:
    tree = ast.parse(Path(path).read_text(encoding="utf-8"))
    module_names = bound_names(tree.body)
    seams = [
        scan(node, module_names)
        for node in ast.walk(tree)
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
    ]
    seams.sort(key=lambda s: (-s.score, s.lineno))
    print(f"{'score':>5}  {'function':<26} writes alias reads calls")
    for seam in seams:
        print(
            f"{seam.score:>5}  {seam.qualname:<26} "
            f"{len(seam.writes):>6} {len(seam.aliases):>5} {len(seam.reads):>5} {len(seam.calls):>5}"
        )
        if seam.writes:
            print(f"       writes: {', '.join(sorted(seam.writes))}")


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

python seam_census.py billing/legacy.py

score  function                   writes alias reads calls
    8  recompute_totals              1     2     1     0
       writes: LEDGER
    5  apply_fee                     0     1     2     1
    2  format_statement              0     0     2     0

# tests/test_legacy_blast_radius.py
import copy
import importlib

import pytest

TARGET = "billing.legacy"


def snapshot(mod):
    snap = {}
    for name, value in vars(mod).items():
        if name.startswith("__") and name.endswith("__"):
            continue
        try:
            snap[name] = copy.deepcopy(value)
        except Exception:
            snap[name] = repr(value)  # last resort, unstable, allowlist it
    return snap


def changed(before, after):
    keys = set(before) | set(after)
    return sorted(k for k in keys if before.get(k) != after.get(k))


@pytest.mark.parametrize(
    "call, expected",
    [
        (lambda m: m.apply_fee("acct-1", 250), ["LEDGER"]),
        (lambda m: m.format_statement("acct-1"), []),
        (lambda m: m.recompute_totals(), ["LEDGER", "_cache", "RATES"]),
    ],
)
def test_blast_radius_is_frozen(call, expected):
    mod = importlib.import_module(TARGET)
    before = snapshot(mod)
    call(mod)
    assert changed(before, snapshot(mod)) == expected