Cst rewrite

cst_rewrite

libcst write helpers + the import-index cache.

Everything that mutates Python source code lives here. Split into five concerns:

• class flatten (_flatten_class_to_top_level, _flatten_class_child)
• rename (rename_name_in_module, rename_top_level_in_source)
• delete (delete_function_from_source, delete_source_if_empty_tests)
• statement reorder (reorder_module_statements + ast helpers)
• Path(__file__).parents[N] depth patch (patch_file_dunder_depth)
• import management — read-side analysis lives in tests_ast; here we insert, dedupe, backfill, and synthesise missing imports, plus the project-wide import index that keeps backfill O(1) per lookup.

Hybrid I/O: we read with ast (cheap, well-tested for analysis) and write with libcst (source-fidelity: comments, triple-quoted strings, blank lines, quote style all preserved — what ast.unparse silently loses).

backfill_import(module, mapping)

Insert from {mod} import {name} for each (name → mod) in mapping.

The new imports go at the canonical position — after every from __future__ import … (or at the module top when none exists) and before the first non-import statement. Names already imported in module are skipped, so this is idempotent on already-correct sources.

Source code in packages/axm-audit/src/axm_audit/core/fix/cst_rewrite.py

def backfill_import(module: cst.Module, mapping: dict[str, str]) -> cst.Module:
    """Insert ``from {mod} import {name}`` for each (name → mod) in *mapping*.

    The new imports go at the canonical position — after every
    ``from __future__ import …`` (or at the module top when none exists)
    and before the first non-import statement. Names already imported in
    *module* are skipped, so this is idempotent on already-correct sources.
    """
    if not mapping:
        return module

    new_imports = _build_import_lines(mapping, _existing_import_names(module))
    if not new_imports:
        return module

    body = list(module.body)
    insert_at = _future_import_insert_index(body)
    return module.with_changes(body=body[:insert_at] + new_imports + body[insert_at:])

backfill_missing_imports(source, target, project_path=None)

Copy imports from source into target for names target uses but doesn't define.

Falls back to scanning all test files under project_path if the immediate source doesn't have the import — covers cases where the original import was lost by an earlier move.

Hybrid: analyse with ast (cheap, well-tested), write with libcst so triple-quoted strings, blank lines, and comments in the target file are preserved byte-for-byte.

Source code in packages/axm-audit/src/axm_audit/core/fix/cst_rewrite.py

def backfill_missing_imports(
    source: Path, target: Path, project_path: Path | None = None
) -> list[str]:
    """Copy imports from *source* into *target* for names target uses but doesn't define.

    Falls back to scanning all test files under ``project_path`` if the
    immediate source doesn't have the import — covers cases where the
    original import was lost by an earlier move.

    Hybrid: analyse with ast (cheap, well-tested), write with libcst so
    triple-quoted strings, blank lines, and comments in the target file
    are preserved byte-for-byte.
    """
    recoverable = _resolve_recoverable_imports(source, target, project_path)
    if recoverable is None:
        return []
    unresolved = recoverable.pop(_UNRESOLVED_KEY, None)
    msgs: list[str] = []
    if recoverable:
        top_level_ast, type_checking_ast, msgs = _partition_imports(
            recoverable, source.name
        )
        _write_backfilled_imports(target, top_level_ast, type_checking_ast)
    if isinstance(unresolved, set):
        msgs.extend(
            f"unresolved import for `{name}` in {target.name} "
            "(no donor found; left for manual fix)"
            for name in sorted(unresolved)
        )
    return msgs

dedupe_imports(module)

Public wrapper around :func:_dedupe_imports_cst.

Source code in packages/axm-audit/src/axm_audit/core/fix/cst_rewrite.py

def dedupe_imports(module: cst.Module) -> cst.Module:
    """Public wrapper around :func:`_dedupe_imports_cst`."""
    return _dedupe_imports_cst(module)

delete_function(module, func_name)

Drop top-level function func_name from module.

Neighbouring statements (and their attached blank-line spacing) are preserved by libcst's leading-lines semantics. In-memory counterpart of :func:delete_function_from_source.

Source code in packages/axm-audit/src/axm_audit/core/fix/cst_rewrite.py

def delete_function(module: cst.Module, func_name: str) -> cst.Module:
    """Drop top-level function *func_name* from *module*.

    Neighbouring statements (and their attached blank-line spacing) are
    preserved by libcst's leading-lines semantics. In-memory counterpart
    of :func:`delete_function_from_source`.
    """
    new_body = [
        stmt
        for stmt in module.body
        if not (isinstance(stmt, cst.FunctionDef) and stmt.name.value == func_name)
    ]
    return module.with_changes(body=new_body)

delete_function_from_source(source, func_name)

Remove a top-level FunctionDef from source, preserving formatting.

Source code in packages/axm-audit/src/axm_audit/core/fix/cst_rewrite.py

def delete_function_from_source(source: Path, func_name: str) -> None:
    """Remove a top-level FunctionDef from source, preserving formatting."""
    module = cst_load(source)
    if module is None:
        return
    new_body = [
        stmt
        for stmt in module.body
        if not (isinstance(stmt, cst.FunctionDef) and stmt.name.value == func_name)
    ]
    cst_save(source, module.with_changes(body=new_body))

delete_source_if_empty_tests(source)

git rm the source if no test_* funcs/classes remain.

Source code in packages/axm-audit/src/axm_audit/core/fix/cst_rewrite.py

def delete_source_if_empty_tests(source: Path) -> None:
    """git rm the source if no test_* funcs/classes remain."""
    if not source.exists():
        return
    tree = ast.parse(source.read_text())
    if _walk_test_funcs(tree):
        return
    rc = subprocess.run(
        ["git", "rm", "-q", str(source)],
        capture_output=True,
        text=True,
    )
    if rc.returncode != 0:
        source.unlink()

flatten_class(module, class_name)

Flatten the class_name class into top-level functions.

In-memory variant of :func:_flatten_class_to_top_level. Class-level pytest marks are propagated onto each promoted method; method-level decorators are preserved verbatim; the class docstring is dropped.

Source code in packages/axm-audit/src/axm_audit/core/fix/cst_rewrite.py

def flatten_class(module: cst.Module, class_name: str) -> cst.Module:
    """Flatten the *class_name* class into top-level functions.

    In-memory variant of :func:`_flatten_class_to_top_level`. Class-level
    pytest marks are propagated onto each promoted method; method-level
    decorators are preserved verbatim; the class docstring is dropped.
    """
    new_body: list[cst.BaseStatement] = []
    for stmt in module.body:
        if not (isinstance(stmt, cst.ClassDef) and stmt.name.value == class_name):
            new_body.append(stmt)
            continue
        class_decos = tuple(d for d in stmt.decorators if _is_pytest_mark_decorator(d))
        for child in stmt.body.body:
            promoted = _flatten_class_child(child, class_decos)
            if promoted is not None:
                new_body.append(promoted)
    return module.with_changes(body=new_body)

invalidate_import_index(project_path)

Drop the cached import index for project_path.

Source code in packages/axm-audit/src/axm_audit/core/fix/cst_rewrite.py

def invalidate_import_index(project_path: Path) -> None:
    """Drop the cached import index for *project_path*."""
    _PROJECT_IMPORT_INDEX_CACHE.pop(project_path, None)

patch_file_depth(module, depth_delta=0)

Rewrite Path(__file__).parents[N] literals by depth_delta.

In-memory variant of :func:patch_file_dunder_depth that targets the subscript form only — the chained .parent.parent form is left for the file-level helper. Identity transform when depth_delta is 0 or the pattern is absent.

Source code in packages/axm-audit/src/axm_audit/core/fix/cst_rewrite.py

def patch_file_depth(module: cst.Module, depth_delta: int = 0) -> cst.Module:
    """Rewrite ``Path(__file__).parents[N]`` literals by *depth_delta*.

    In-memory variant of :func:`patch_file_dunder_depth` that targets the
    subscript form only — the chained ``.parent.parent`` form is left for
    the file-level helper. Identity transform when *depth_delta* is 0 or
    the pattern is absent.
    """
    if depth_delta == 0:
        return module

    class _DunderPatcher(cst.CSTTransformer):
        def leave_Subscript(
            self,
            original_node: cst.Subscript,
            updated_node: cst.Subscript,
        ) -> cst.BaseExpression:
            value = updated_node.value
            if not isinstance(value, cst.Attribute):
                return updated_node
            if value.attr.value != "parents":
                return updated_node
            if not _is_file_dunder_chain(value.value):
                return updated_node
            slices = updated_node.slice
            if len(slices) != 1:
                return updated_node
            elt = slices[0].slice
            if not isinstance(elt, cst.Index):
                return updated_node
            n_node = elt.value
            if not isinstance(n_node, cst.Integer):
                return updated_node
            new_n = int(n_node.value) + depth_delta
            if new_n <= 0:
                return updated_node
            return updated_node.with_changes(
                slice=[
                    cst.SubscriptElement(
                        slice=cst.Index(value=cst.Integer(value=str(new_n)))
                    )
                ]
            )

    result = module.visit(_DunderPatcher())
    assert isinstance(result, cst.Module)
    return result

patch_file_dunder_depth(file, depth_delta)

Rewrite Path(__file__).parents[N] / .parent.parent... after a move.

When a file is relocated by depth_delta directory levels (depth_delta = target_depth - source_depth; negative if moved closer to project root, positive if moved deeper), any constant of the form Path(__file__).parents[N] or Path(__file__).parent.parent... will resolve to a different ancestor unless N is adjusted. We compute the new N so the constant continues to resolve to the same directory it did before the move:

Text Only

N_new = N_old + depth_delta

Reasoning: a file at depth D has parents[N] at depth D - N - 1 from project root. Moving the file to depth D' means parents[N'] is at D' - N' - 1. For these to be equal, N' = N + (D' - D) = N + depth_delta.

Two surface forms supported (in order of preference, since some files mix them):

• Subscript: Path(__file__).parents[N] (with optional .resolve()). N is decremented by depth_delta.
• Chained: Path(__file__).parent.parent[.parent]* (with optional .resolve()). The number of .parent accessors is reduced by depth_delta.

If the resulting N would be <= 0, we leave the constant alone and emit a warning — the file was moved too close to root for the resolution to be expressible, indicating the relocate is suspect.

Source code in packages/axm-audit/src/axm_audit/core/fix/cst_rewrite.py

def patch_file_dunder_depth(
    file: Path,
    depth_delta: int,
) -> list[str]:
    """Rewrite ``Path(__file__).parents[N]`` / ``.parent.parent...`` after a move.

    When a file is relocated by *depth_delta* directory levels
    (``depth_delta = target_depth - source_depth``; negative if moved
    closer to project root, positive if moved deeper), any constant of
    the form ``Path(__file__).parents[N]`` or
    ``Path(__file__).parent.parent...`` will resolve to a different
    ancestor unless ``N`` is adjusted. We compute the new ``N`` so the
    constant continues to resolve to the *same* directory it did before
    the move:

        N_new = N_old + depth_delta

    Reasoning: a file at depth ``D`` has ``parents[N]`` at depth
    ``D - N - 1`` from project root. Moving the file to depth ``D'``
    means ``parents[N']`` is at ``D' - N' - 1``. For these to be equal,
    ``N' = N + (D' - D)`` = ``N + depth_delta``.

    Two surface forms supported (in order of preference, since some
    files mix them):

      * Subscript: ``Path(__file__).parents[N]`` (with optional
        ``.resolve()``). N is decremented by ``depth_delta``.
      * Chained: ``Path(__file__).parent.parent[.parent]*`` (with
        optional ``.resolve()``). The number of ``.parent`` accessors
        is reduced by ``depth_delta``.

    If the resulting N would be ``<= 0``, we leave the constant alone
    and emit a warning — the file was moved too close to root for the
    resolution to be expressible, indicating the relocate is suspect.
    """
    if depth_delta == 0 or not file.exists():
        return []
    module = cst_load(file)
    if module is None:
        return []
    ctx = _DepthPatchCtx(file=file, depth_delta=depth_delta)
    new_module = module.visit(_DunderSubscriptPatcher(ctx))
    assert isinstance(new_module, cst.Module)
    # Collect ids AFTER _DunderSubscriptPatcher: libcst rebuilds nodes
    # during a visit even when no transformation is returned, so ids
    # captured on `module` would not match the nodes inside `new_module`.
    collector = _CollectChainChildren()
    new_module.visit(collector)
    new_module = new_module.visit(_PatchChainOnce(ctx, collector.child_ids))
    assert isinstance(new_module, cst.Module)
    if new_module.code != module.code:
        cst_save(file, new_module)
    return ctx.msgs

rename_function(module, old_name, new_name)

Rename top-level function old_name to new_name across module.

Updates the FunctionDef itself, any Name reference, and any string-literal argument (e.g. pytest.mark.parametrize("old", …)) that matches old_name. In-memory counterpart of :func:rename_name_in_module.

Source code in packages/axm-audit/src/axm_audit/core/fix/cst_rewrite.py

def rename_function(module: cst.Module, old_name: str, new_name: str) -> cst.Module:
    """Rename top-level function *old_name* to *new_name* across *module*.

    Updates the ``FunctionDef`` itself, any ``Name`` reference, and any
    string-literal argument (e.g. ``pytest.mark.parametrize("old", …)``)
    that matches *old_name*. In-memory counterpart of
    :func:`rename_name_in_module`.
    """
    mapping = {old_name: new_name}

    class _Renamer(cst.CSTTransformer):
        def leave_Name(
            self, original_node: cst.Name, updated_node: cst.Name
        ) -> cst.BaseExpression:
            if updated_node.value in mapping:
                return updated_node.with_changes(value=mapping[updated_node.value])
            return updated_node

        def leave_FunctionDef(
            self,
            original_node: cst.FunctionDef,
            updated_node: cst.FunctionDef,
        ) -> cst.BaseStatement:
            if updated_node.name.value in mapping:
                return updated_node.with_changes(
                    name=cst.Name(value=mapping[updated_node.name.value])
                )
            return updated_node

        def leave_SimpleString(
            self,
            original_node: cst.SimpleString,
            updated_node: cst.SimpleString,
        ) -> cst.BaseExpression:
            raw = updated_node.value
            if len(raw) < 2 or raw[0] not in {'"', "'"}:
                return updated_node
            inner = raw[1:-1]
            if inner in mapping:
                return updated_node.with_changes(
                    value=f"{raw[0]}{mapping[inner]}{raw[0]}"
                )
            return updated_node

    result = module.visit(_Renamer())
    assert isinstance(result, cst.Module)
    return result

rename_name_in_module(path, old_to_new)

Rename every occurrence of name X across module path (def + refs).

Renames at three sites simultaneously

• the cst.FunctionDef / cst.ClassDef definition itself,
• every cst.Name reference in the module body,
• marker-argument string literals like @pytest.mark.usefixtures("X") so usefixtures still resolves after the rename.

Preserves formatting via libcst. Unlike rename_top_level_in_source (which only renames the def header — needed for cross-file move collisions), this rewrites references too — needed when source helpers get renamed to avoid colliding with target's same-named helpers.

Source code in packages/axm-audit/src/axm_audit/core/fix/cst_rewrite.py

def rename_name_in_module(path: Path, old_to_new: dict[str, str]) -> None:
    """Rename every occurrence of name X across module *path* (def + refs).

    Renames at three sites simultaneously:
      * the ``cst.FunctionDef`` / ``cst.ClassDef`` definition itself,
      * every ``cst.Name`` reference in the module body,
      * marker-argument string literals like
        ``@pytest.mark.usefixtures("X")`` so usefixtures still resolves
        after the rename.

    Preserves formatting via libcst. Unlike
    ``rename_top_level_in_source`` (which only renames the def header
    — needed for cross-file move collisions), this rewrites references
    too — needed when source helpers get renamed to avoid colliding with
    target's same-named helpers.
    """
    if not old_to_new:
        return
    module = cst_load(path)
    if module is None:
        return

    class _Renamer(cst.CSTTransformer):
        def __init__(self, mapping: dict[str, str]) -> None:
            self.mapping = mapping

        def leave_Name(
            self, original_node: cst.Name, updated_node: cst.Name
        ) -> cst.BaseExpression:
            if updated_node.value in self.mapping:
                return updated_node.with_changes(value=self.mapping[updated_node.value])
            return updated_node

        def leave_FunctionDef(
            self,
            original_node: cst.FunctionDef,
            updated_node: cst.FunctionDef,
        ) -> cst.BaseStatement:
            if updated_node.name.value in self.mapping:
                return updated_node.with_changes(
                    name=cst.Name(value=self.mapping[updated_node.name.value])
                )
            return updated_node

        def leave_ClassDef(
            self,
            original_node: cst.ClassDef,
            updated_node: cst.ClassDef,
        ) -> cst.BaseStatement:
            if updated_node.name.value in self.mapping:
                return updated_node.with_changes(
                    name=cst.Name(value=self.mapping[updated_node.name.value])
                )
            return updated_node

        def leave_SimpleString(
            self,
            original_node: cst.SimpleString,
            updated_node: cst.SimpleString,
        ) -> cst.BaseExpression:
            raw = updated_node.value
            if len(raw) < 2:
                return updated_node
            quote = raw[0]
            if quote not in {'"', "'"}:
                return updated_node
            inner = raw[1:-1]
            if inner in self.mapping:
                return updated_node.with_changes(
                    value=f"{quote}{self.mapping[inner]}{quote}"
                )
            return updated_node

    new_module = module.visit(_Renamer(old_to_new))
    assert isinstance(new_module, cst.Module)
    cst_save(path, new_module)

rename_top_level_in_source(source, old_to_new)

Rename top-level FunctionDef / ClassDef in source, preserving formatting.

Workaround for axm-anvil's rename= parameter, which validates target absence under the ORIGINAL name before applying the rename — so it cannot resolve cross-file collisions on its own. By renaming in source first, we hand anvil a clean conflict-free move.

Source code in packages/axm-audit/src/axm_audit/core/fix/cst_rewrite.py

def rename_top_level_in_source(source: Path, old_to_new: dict[str, str]) -> None:
    """Rename top-level FunctionDef / ClassDef in *source*, preserving formatting.

    Workaround for axm-anvil's ``rename=`` parameter, which validates
    target absence under the ORIGINAL name before applying the rename —
    so it cannot resolve cross-file collisions on its own. By renaming in
    source first, we hand anvil a clean conflict-free move.
    """
    if not old_to_new:
        return
    module = cst_load(source)
    if module is None:
        return
    new_body = []
    for stmt in module.body:
        if (
            isinstance(stmt, cst.FunctionDef | cst.ClassDef)
            and stmt.name.value in old_to_new
        ):
            stmt = stmt.with_changes(name=cst.Name(value=old_to_new[stmt.name.value]))
        new_body.append(stmt)
    cst_save(source, module.with_changes(body=new_body))

reorder_module_statements(path)

Reorder a module's top-level statements so definitions precede uses.

After SPLIT/MERGE/FLATTEN, axm-anvil can leave statements in an order that breaks Python's module-execution semantics: * _skip_no_tools = pytest.mark.skipif(_tools_available()) before def _tools_available() → NameError at import. * @_skip_no_tools decorator on a class, before the assign that defines _skip_no_tools → NameError.

Strategy: stable topological sort. Imports stay first (they have no intra-module deps). For the rest, each statement is placed after the last statement that defines a name it references at module-execution time. References inside function bodies do NOT count — they're deferred. Order is preserved within independent groups.

Implementation note: we parse twice — once with libcst (the source of truth for formatting; what we'll write back) and once with ast (for cheap defines/references analysis). The libcst statements are reordered by index, not rebuilt, so triple-quoted strings, comments, and blank-line spacing all survive intact.

Idempotent.

Source code in packages/axm-audit/src/axm_audit/core/fix/cst_rewrite.py

def reorder_module_statements(path: Path) -> None:
    """Reorder a module's top-level statements so definitions precede uses.

    After SPLIT/MERGE/FLATTEN, axm-anvil can leave statements in an order
    that breaks Python's module-execution semantics:
      * ``_skip_no_tools = pytest.mark.skipif(_tools_available())`` before
        ``def _tools_available()`` → NameError at import.
      * ``@_skip_no_tools`` decorator on a class, before the assign that
        defines ``_skip_no_tools`` → NameError.

    Strategy: stable topological sort. Imports stay first (they have no
    intra-module deps). For the rest, each statement is placed after the
    last statement that defines a name it references at module-execution
    time. References inside function bodies do NOT count — they're
    deferred. Order is preserved within independent groups.

    Implementation note: we parse twice — once with libcst (the source
    of truth for formatting; what we'll write back) and once with ast
    (for cheap defines/references analysis). The libcst statements are
    reordered by index, not rebuilt, so triple-quoted strings, comments,
    and blank-line spacing all survive intact.

    Idempotent.
    """
    loaded = _load_cst_ast_pair(path)
    if loaded is None:
        return
    cst_module, body_ast = loaded
    body_cst = list(cst_module.body)

    docstring_idx = _find_docstring_idx(body_ast)
    head_idx, rest_idx = _partition_head_rest(body_ast, docstring_idx)
    earliest, needs_change = _compute_earliest(body_ast, head_idx, rest_idx)

    docstring_misplaced = docstring_idx is not None and docstring_idx != 0
    if not needs_change and not docstring_misplaced:
        return

    new_body_cst = _build_reordered_body(
        body_cst, head_idx, rest_idx, earliest, docstring_idx
    )
    new_text = cst_module.with_changes(body=new_body_cst).code
    if new_text != cst_module.code:
        path.write_text(new_text)

resolve_import_for_symbol(project_path, symbol)

Return the import statement that brings symbol into scope, or None.

Builds (and caches in _PROJECT_IMPORT_INDEX_CACHE) a project-wide index of top-level FunctionDef / AsyncFunctionDef / ClassDef definitions across every .py file under project_path. Drop the cache via :func:invalidate_import_index after mutating the file tree so the next call rebuilds.

Source code in packages/axm-audit/src/axm_audit/core/fix/cst_rewrite.py

def resolve_import_for_symbol(
    project_path: Path, symbol: str
) -> tuple[ast.stmt, ast.stmt | None] | None:
    """Return the import statement that brings *symbol* into scope, or ``None``.

    Builds (and caches in ``_PROJECT_IMPORT_INDEX_CACHE``) a project-wide
    index of top-level FunctionDef / AsyncFunctionDef / ClassDef
    definitions across every ``.py`` file under *project_path*. Drop the
    cache via :func:`invalidate_import_index` after mutating the file
    tree so the next call rebuilds.
    """
    if project_path not in _PROJECT_IMPORT_INDEX_CACHE:
        _PROJECT_IMPORT_INDEX_CACHE[project_path] = _build_project_symbol_index(
            project_path
        )
    return _PROJECT_IMPORT_INDEX_CACHE[project_path].get(symbol)