"""muse code rename — symbol-aware rename across definition, imports, and call sites. Unlike ``sed -i 's/old/new/g'``, ``muse code rename`` operates at the AST level — it understands Python's parse tree and only rewrites the identifier token at each relevant site: - **Definition** — the ``def``/``class`` name token, scoped to the correct class body for methods. String literals and comments that happen to contain the old name are never touched. - **Import sites** — every ``from module import old_name`` (and ``import old_name``) in the repository. - **Call / reference sites** — every bare ``old_name(...)`` usage (``ast.Name``) and every attribute access ``obj.old_name(...)`` (``ast.Attribute``). Each edit is character-precise: the command records the exact (line, col) range and replaces only that token, leaving the surrounding text intact. The command is **dry-run by default** — pass ``--yes`` (or ``-y``) to write. Agents should use ``--json --yes`` for fully automated pipelines. Usage:: # Preview all changes muse code rename billing.py::compute_total compute_invoice_total # Apply without prompt muse code rename billing.py::compute_total compute_invoice_total --yes # Methods (scoped to the class body) muse code rename billing.py::Invoice.compute_total compute_invoice_total -y # Only rename the definition — leave call sites for manual review muse code rename billing.py::compute_total compute_invoice_total \\ --scope definition --yes # Machine-readable output for agents muse code rename billing.py::compute_total compute_invoice_total --json # Larger codebases — lift the file cap muse code rename billing.py::compute_total compute_invoice_total \\ --max-files 200 --yes Output (human):: Renaming billing.py::compute_total → compute_invoice_total Definition billing.py:4 def compute_total(self, items): Import sites tests/test_billing.py:2 from billing import compute_total Call / reference sites services/order.py:15 result = compute_total(items) tests/test_billing.py:8 total = compute_total([1, 2, 3]) 4 edit site(s) across 3 file(s) Apply changes? [y/N] """ from __future__ import annotations import ast import argparse import json import logging import pathlib import re import sys from collections import defaultdict from typing import TypedDict from muse.core.errors import ExitCode from muse.core.repo import read_repo_id, require_repo from muse.core.store import ( get_commit_snapshot_manifest, read_current_branch, resolve_commit_ref, Manifest, ) from muse.plugins.code._query import language_of from muse.core.validation import clamp_int, MAX_AST_BYTES, sanitize_display type _EditSiteMap = dict[str, list["_EditSite"]] logger = logging.getLogger(__name__) # ── Constants ────────────────────────────────────────────────────────────────── _MAX_SYMBOL_ADDR_LEN = 500 _MAX_NAME_LEN = 200 _DEFAULT_MAX_FILES = 100 # Valid Python identifier (post-normalisation). _IDENT_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$") # Dunder names require --force to rename. _DUNDER_RE = re.compile(r"^__[A-Za-z_][A-Za-z0-9_]*__$") _SCOPE_CHOICES = ("all", "definition", "imports", "callsites") # ── TypedDicts ───────────────────────────────────────────────────────────────── class _EditSite(TypedDict): """A single token replacement site.""" file: str line: int # 1-indexed col_start: int # 0-indexed, inclusive col_end: int # 0-indexed, exclusive kind: str # "definition" | "import" | "reference" context: str # full source line for display / JSON class _RenameResult(TypedDict): """Machine-readable summary of a rename operation.""" from_address: str to_address: str from_name: str to_name: str scope: str dry_run: bool files_to_modify: list[str] total_edit_sites: int edit_sites: list[_EditSite] warnings: list[str] # ── Helpers ──────────────────────────────────────────────────────────────────── def _validate_identifier(name: str, force: bool) -> str | None: """Return an error string if *name* is not a valid rename target, else None.""" if not name: return "New name must not be empty." if len(name) > _MAX_NAME_LEN: return f"New name too long (max {_MAX_NAME_LEN} chars)." if not _IDENT_RE.match(name): return f"'{name}' is not a valid Python identifier." if _DUNDER_RE.match(name) and not force: return ( f"'{name}' is a dunder name. Pass --force to rename it anyway." ) return None def _parse_address(address: str) -> tuple[str, list[str]] | None: """Split 'billing.py::Invoice.method' into ('billing.py', ['Invoice', 'method']). Returns None if the address is malformed. """ if "::" not in address: return None file_part, _, symbol_part = address.partition("::") if not file_part or not symbol_part: return None name_parts = symbol_part.split(".") if any(not p for p in name_parts): return None return file_part, name_parts def _line(lines: list[str], lineno: int) -> str: """Return the source line (1-indexed) or empty string if out of range.""" return lines[lineno - 1] if 1 <= lineno <= len(lines) else "" # ── Definition finder ────────────────────────────────────────────────────────── def _find_definition_site( source: str, file_path: str, name_parts: list[str], ) -> _EditSite | None: """Locate the AST position of the symbol's def/class token. Respects class scope: ``billing.py::Invoice.compute_total`` only matches the ``def compute_total`` *inside* the ``Invoice`` class body. """ try: if len(source) > MAX_AST_BYTES: return None tree = ast.parse(source, filename=file_path) except SyntaxError: return None lines = source.splitlines() target = name_parts[-1] def _site( node: ast.FunctionDef | ast.AsyncFunctionDef | ast.ClassDef, ) -> _EditSite: if isinstance(node, ast.AsyncFunctionDef): prefix_len = len("async def ") elif isinstance(node, ast.FunctionDef): prefix_len = len("def ") else: prefix_len = len("class ") col_start = node.col_offset + prefix_len col_end = col_start + len(target) return _EditSite( file=file_path, line=node.lineno, col_start=col_start, col_end=col_end, kind="definition", context=_line(lines, node.lineno), ) if len(name_parts) == 1: # Module-level function or class. for node in tree.body: if ( isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) and node.name == target ): return _site(node) if isinstance(node, ast.ClassDef) and node.name == target: return _site(node) elif len(name_parts) >= 2: # Method (or nested class) — walk into the correct parent class. class_name = name_parts[-2] for top in tree.body: if isinstance(top, ast.ClassDef) and top.name == class_name: for item in ast.walk(top): if ( isinstance(item, (ast.FunctionDef, ast.AsyncFunctionDef)) and item.name == target ): return _site(item) if isinstance(item, ast.ClassDef) and item.name == target: return _site(item) return None # ── Reference / import finder ────────────────────────────────────────────────── def _find_reference_sites( source: str, file_path: str, old_name: str, include_imports: bool, include_callsites: bool, ) -> list[_EditSite]: """Find all import and call / reference sites of *old_name* in *source*. Does NOT return the definition site — call :func:`_find_definition_site` for that separately. """ try: if len(source) > MAX_AST_BYTES: return [] tree = ast.parse(source, filename=file_path) except SyntaxError: return [] lines = source.splitlines() sites: list[_EditSite] = [] pattern = re.compile(r"\b" + re.escape(old_name) + r"\b") for node in ast.walk(tree): # ── Import sites ───────────────────────────────────────────────────── if include_imports and isinstance(node, ast.ImportFrom): for alias in node.names: if alias.name == old_name: # Python's AST does not expose per-alias column offsets, so # we use a word-boundary regex against the specific import # line. This is safe because we already know old_name # appears here as an imported name. import_line = _line(lines, node.lineno) for m in pattern.finditer(import_line): sites.append(_EditSite( file=file_path, line=node.lineno, col_start=m.start(), col_end=m.end(), kind="import", context=import_line, )) if include_imports and isinstance(node, ast.Import): for alias in node.names: if alias.name == old_name or alias.name.split(".")[-1] == old_name: import_line = _line(lines, node.lineno) for m in pattern.finditer(import_line): sites.append(_EditSite( file=file_path, line=node.lineno, col_start=m.start(), col_end=m.end(), kind="import", context=import_line, )) if not include_callsites: continue # ── Bare name references (ast.Name) ────────────────────────────────── if isinstance(node, ast.Name) and node.id == old_name: sites.append(_EditSite( file=file_path, line=node.lineno, col_start=node.col_offset, col_end=node.col_offset + len(old_name), kind="reference", context=_line(lines, node.lineno), )) # ── Attribute access (ast.Attribute) — obj.old_name ────────────────── elif isinstance(node, ast.Attribute) and node.attr == old_name: end_line: int = node.end_lineno if node.end_lineno is not None else node.lineno end_col: int = node.end_col_offset if node.end_col_offset is not None else 0 col_start = end_col - len(old_name) if col_start < 0: continue # defensive: malformed node sites.append(_EditSite( file=file_path, line=end_line, col_start=col_start, col_end=end_col, kind="reference", context=_line(lines, end_line), )) return sites # ── Edit applicator ──────────────────────────────────────────────────────────── def _apply_edits(source: str, sites: list[_EditSite], new_name: str) -> str: """Apply all edit sites to *source*, replacing each token with *new_name*. Edits are applied right-to-left within each line so that column offsets remain valid for subsequent replacements on the same line. """ lines = source.splitlines(keepends=True) by_line: dict[int, list[_EditSite]] = defaultdict(list) for site in sites: by_line[site["line"]].append(site) result: list[str] = [] for i, raw_line in enumerate(lines): lineno = i + 1 if lineno not in by_line: result.append(raw_line) continue # Strip the trailing newline, apply edits, then restore it. stripped = raw_line.rstrip("\r\n") tail = raw_line[len(stripped):] # Sort descending by col_start so right-to-left replacement works. edits = sorted(by_line[lineno], key=lambda s: -s["col_start"]) for edit in edits: cs, ce = edit["col_start"], edit["col_end"] # Defensive bounds check. cs = max(0, min(cs, len(stripped))) ce = max(cs, min(ce, len(stripped))) stripped = stripped[:cs] + new_name + stripped[ce:] result.append(stripped + tail) return "".join(result) # ── Deduplication ────────────────────────────────────────────────────────────── def _dedup(sites: list[_EditSite]) -> list[_EditSite]: """Remove duplicate sites (same file, line, col_start).""" seen: set[tuple[int, int]] = set() out: list[_EditSite] = [] for s in sites: key = (s["line"], s["col_start"]) if key not in seen: seen.add(key) out.append(s) return out # ── CLI registration ─────────────────────────────────────────────────────────── def register( subparsers: "argparse._SubParsersAction[argparse.ArgumentParser]", ) -> None: """Register the rename subcommand.""" parser = subparsers.add_parser( "rename", help=( "Rename a symbol across its definition, all import sites, " "and all call sites — AST-level, never touches string literals." ), description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter, ) parser.add_argument( "address", metavar="ADDRESS", help=( "Symbol to rename (e.g. billing.py::compute_total or " "billing.py::Invoice.compute_total)." ), ) parser.add_argument( "new_name", metavar="NEW_NAME", help="New bare identifier (e.g. compute_invoice_total).", ) parser.add_argument( "--scope", default="all", choices=_SCOPE_CHOICES, metavar="SCOPE", help=( f"What to rename: {', '.join(_SCOPE_CHOICES)} (default: all). " "'definition' renames only the def/class token; 'imports' updates " "only import statements; 'callsites' updates only call / reference " "sites. 'all' does all three." ), ) parser.add_argument( "--yes", "-y", dest="yes", action="store_true", help="Apply without an interactive confirmation prompt (for agents).", ) parser.add_argument( "--dry-run", "-n", dest="dry_run", action="store_true", help="Show what would change without writing any files.", ) parser.add_argument( "--json", dest="as_json", action="store_true", help=( "Emit a machine-readable JSON summary of every edit site. " "Combine with --yes to apply in one step." ), ) parser.add_argument( "--max-files", type=int, default=_DEFAULT_MAX_FILES, metavar="N", dest="max_files", help=( f"Maximum number of Python files to scan for references " f"(default: {_DEFAULT_MAX_FILES}). Increase for large repos." ), ) parser.add_argument( "--force", action="store_true", help="Allow renaming dunder methods (__init__, __str__, …).", ) parser.set_defaults(func=run) # ── Main logic ───────────────────────────────────────────────────────────────── def run(args: argparse.Namespace) -> None: """Rename a symbol across its definition, imports, and call sites. Operates exclusively on the working tree (HEAD snapshot for file discovery, disk files for content). After renaming, run ``muse status`` then ``muse commit`` as normal. """ address: str = args.address new_name: str = args.new_name scope: str = args.scope dry_run: bool = args.dry_run yes: bool = args.yes as_json: bool = args.as_json max_files: int = clamp_int(args.max_files, 1, 10000, 'max_files') force: bool = args.force # ── Input validation ────────────────────────────────────────────────────── if len(address) > _MAX_SYMBOL_ADDR_LEN: print( f"❌ ADDRESS too long (max {_MAX_SYMBOL_ADDR_LEN} chars).", file=sys.stderr, ) raise SystemExit(ExitCode.USER_ERROR) if "::" not in address: print( "❌ ADDRESS must contain '::' (e.g. billing.py::compute_total).", file=sys.stderr, ) raise SystemExit(ExitCode.USER_ERROR) if max_files < 1: print("❌ --max-files must be >= 1.", file=sys.stderr) raise SystemExit(ExitCode.USER_ERROR) id_error = _validate_identifier(new_name, force) if id_error: print(f"❌ {id_error}", file=sys.stderr) raise SystemExit(ExitCode.USER_ERROR) parsed = _parse_address(address) if parsed is None: print(f"❌ Cannot parse address '{sanitize_display(address)}'.", file=sys.stderr) raise SystemExit(ExitCode.USER_ERROR) file_path, name_parts = parsed old_name = name_parts[-1] if old_name == new_name: print("❌ New name is identical to the old name.", file=sys.stderr) raise SystemExit(ExitCode.USER_ERROR) # ── Locate repo ─────────────────────────────────────────────────────────── root = require_repo() repo_id = read_repo_id(root) branch = read_current_branch(root) head = resolve_commit_ref(root, repo_id, branch, None) if head is None: print("❌ HEAD commit not found.", file=sys.stderr) raise SystemExit(ExitCode.USER_ERROR) manifest: Manifest = get_commit_snapshot_manifest(root, head.commit_id) or {} # ── Security: path containment ──────────────────────────────────────────── # Reject addresses that would escape the repository root. target_abs = (root / file_path).resolve() try: target_abs.relative_to(root.resolve()) except ValueError: print( f"❌ Path '{file_path}' escapes the repository root.", file=sys.stderr, ) raise SystemExit(ExitCode.USER_ERROR) if not target_abs.is_file(): print( f"❌ '{file_path}' not found in the working tree.", file=sys.stderr, ) raise SystemExit(ExitCode.USER_ERROR) # ── Scope flags ─────────────────────────────────────────────────────────── include_definition = scope in ("all", "definition") include_imports = scope in ("all", "imports") include_callsites = scope in ("all", "callsites") # ── Collect definition site ─────────────────────────────────────────────── all_sites: _EditSiteMap = {} if include_definition: def_source = target_abs.read_text(encoding="utf-8") def_site = _find_definition_site(def_source, file_path, name_parts) if def_site is None: print( f"❌ Symbol '{old_name}' not found in '{file_path}'.", file=sys.stderr, ) print( f" Hint: muse code symbols --file {file_path}", file=sys.stderr, ) raise SystemExit(ExitCode.USER_ERROR) all_sites[file_path] = [def_site] # ── Collect reference / import sites ────────────────────────────────────── run_warnings: list[str] = [] if include_imports or include_callsites: # Gather all Python files in the snapshot. python_files = sorted( fp for fp in manifest if language_of(fp) == "Python" ) if len(python_files) > max_files: msg = ( f"{len(python_files)} Python files found; scanning only the " f"first {max_files} (pass --max-files {len(python_files)} to scan all)." ) run_warnings.append(msg) if not as_json: print(f"⚠️ {msg}", file=sys.stderr) python_files = python_files[:max_files] # The definition file is always included regardless of ordering. if file_path not in python_files: python_files = [file_path, *python_files] for fp in python_files: disk_path = root / fp if not disk_path.is_file(): continue source = disk_path.read_text(encoding="utf-8") ref_sites = _find_reference_sites( source, fp, old_name, include_imports=include_imports, include_callsites=include_callsites, ) # Remove the definition token from reference results — it is already # captured by _find_definition_site and must not be double-applied. if fp == file_path and file_path in all_sites: def_key = ( all_sites[file_path][0]["line"], all_sites[file_path][0]["col_start"], ) ref_sites = [ s for s in ref_sites if (s["line"], s["col_start"]) != def_key ] if ref_sites: all_sites.setdefault(fp, []).extend(ref_sites) # ── Deduplicate ─────────────────────────────────────────────────────────── for fp in all_sites: all_sites[fp] = _dedup(all_sites[fp]) total_sites = sum(len(v) for v in all_sites.values()) files_affected = sorted(all_sites) # ── Build new address ───────────────────────────────────────────────────── new_parts = name_parts[:-1] + [new_name] new_address = file_path + "::" + ".".join(new_parts) # ── Method rename warning ───────────────────────────────────────────────── # Attribute renames (obj.old_name) are potentially ambiguous when multiple # classes define a method with the same name. Warn in human mode. is_method = len(name_parts) >= 2 attr_sites_found = any( s["kind"] == "reference" for sites in all_sites.values() for s in sites ) # ── JSON output ─────────────────────────────────────────────────────────── if as_json: flat: list[_EditSite] = [] for fp in files_affected: flat.extend(sorted(all_sites[fp], key=lambda s: s["line"])) result: _RenameResult = { "from_address": address, "to_address": new_address, "from_name": old_name, "to_name": new_name, "scope": scope, "dry_run": dry_run, "files_to_modify": files_affected if not dry_run else [], "total_edit_sites": total_sites, "edit_sites": flat, "warnings": run_warnings, } print(json.dumps(result, indent=2)) if dry_run: return # Fall through to apply. # ── Human output: preview ───────────────────────────────────────────────── if not as_json: print(f"\nRenaming {sanitize_display(address)} → {sanitize_display(new_name)}\n") for fp in files_affected: sites = sorted(all_sites[fp], key=lambda s: s["line"]) defs = [s for s in sites if s["kind"] == "definition"] imports = [s for s in sites if s["kind"] == "import"] refs = [s for s in sites if s["kind"] == "reference"] if defs: print(" Definition") for s in defs: print(f" {sanitize_display(fp)}:{s['line']} {sanitize_display(s['context'].rstrip())}") if imports: print(" Import sites") for s in imports: print(f" {sanitize_display(fp)}:{s['line']} {sanitize_display(s['context'].rstrip())}") if refs: print(" Call / reference sites") for s in refs: print(f" {sanitize_display(fp)}:{s['line']} {sanitize_display(s['context'].rstrip())}") print(f"\n {total_sites} edit site(s) across {len(files_affected)} file(s)") if is_method and attr_sites_found: print( "\n ⚠️ Method rename: attribute call sites (obj.old_name) may " "include false positives\n" " if other classes define a method with the same name.\n" " Review with --dry-run, then pass --scope definition if needed." ) if dry_run: print("\n (dry run — no files written)") return if not yes: try: answer = input("\nApply changes? [y/N] ").strip().lower() except (EOFError, KeyboardInterrupt): print("\nAborted.", file=sys.stderr) raise SystemExit(ExitCode.USER_ERROR) if answer not in ("y", "yes"): print("Aborted.", file=sys.stderr) raise SystemExit(ExitCode.USER_ERROR) # ── Apply edits ─────────────────────────────────────────────────────────── for fp in files_affected: disk_path = root / fp if not disk_path.is_file(): logger.warning("⚠️ Skipping %s — not found on disk", fp) continue original = disk_path.read_text(encoding="utf-8") modified = _apply_edits(original, all_sites[fp], new_name) if modified != original: disk_path.write_text(modified, encoding="utf-8") if not as_json: print(f" ✅ {sanitize_display(fp)}") if not as_json: print( f"\n{total_sites} rename(s) applied across {len(files_affected)} file(s)." ) print("Run `muse status` to review, then `muse commit`.")