add_json_alias.py
python
sha256:1900655993c83c4107067375548a7be823e471d2515830842f1a12cba4bd3cdf
fix: unified object store migration — idempotent writes, JS…
Sonnet 4.6
minor
⚠ breaking
28 days ago
| 1 | """Add --json alias to every register() that uses --format/dest='fmt'. |
| 2 | |
| 3 | Uses muse's own semantic tooling: |
| 4 | - muse code query → find every register() in cli/commands |
| 5 | - muse code cat → read the exact source of each one |
| 6 | - muse code patch → surgically replace only the functions that need updating |
| 7 | |
| 8 | No regex, no AST manipulation, no risk of touching surrounding code. |
| 9 | """ |
| 10 | from __future__ import annotations |
| 11 | |
| 12 | import json |
| 13 | import re |
| 14 | import subprocess |
| 15 | import sys |
| 16 | import tempfile |
| 17 | import pathlib |
| 18 | |
| 19 | |
| 20 | # Match any var.add_argument("--format" ...) call that includes dest="fmt" |
| 21 | # but is NOT already followed by a --json counterpart for that same var. |
| 22 | _FORMAT_RE = re.compile( |
| 23 | r'(?m)^(?P<indent>[ \t]*)(?P<var>\w+)\.add_argument\(' |
| 24 | r'[^)]*"--format"[^)]*dest="fmt"[^)]*\)', |
| 25 | ) |
| 26 | |
| 27 | _JSON_ALIAS_RE = re.compile(r'"--json"[^)]*dest="fmt"') |
| 28 | |
| 29 | |
| 30 | def _muse(*args: str) -> str: |
| 31 | """Run a muse command and return stdout. Raises on non-zero exit.""" |
| 32 | result = subprocess.run( |
| 33 | ["muse", *args], |
| 34 | capture_output=True, text=True, |
| 35 | ) |
| 36 | if result.returncode not in (0, 32): # 32 = muse query success code |
| 37 | raise RuntimeError( |
| 38 | f"muse {' '.join(args)} failed ({result.returncode}):\n{result.stderr}" |
| 39 | ) |
| 40 | return result.stdout |
| 41 | |
| 42 | |
| 43 | def _vars_with_json(source: str) -> set[str]: |
| 44 | """Return parser variable names that already have a --json add_argument call. |
| 45 | |
| 46 | Handles both single-line and multi-line calls correctly by tracking paren depth. |
| 47 | """ |
| 48 | found: set[str] = set() |
| 49 | lines = source.splitlines(keepends=True) |
| 50 | i = 0 |
| 51 | while i < len(lines): |
| 52 | m = re.match(r'^[ \t]*(?P<var>\w+)\.add_argument\(', lines[i]) |
| 53 | if m: |
| 54 | var = m.group("var") |
| 55 | # Collect entire call body. |
| 56 | call_lines = [lines[i]] |
| 57 | depth = lines[i].count("(") - lines[i].count(")") |
| 58 | j = i + 1 |
| 59 | while depth > 0 and j < len(lines): |
| 60 | call_lines.append(lines[j]) |
| 61 | depth += lines[j].count("(") - lines[j].count(")") |
| 62 | j += 1 |
| 63 | call_text = "".join(call_lines) |
| 64 | if '"--json"' in call_text: |
| 65 | found.add(var) |
| 66 | i = j |
| 67 | continue |
| 68 | i += 1 |
| 69 | return found |
| 70 | |
| 71 | |
| 72 | def _add_json_aliases(source: str) -> tuple[str, int]: |
| 73 | """Return (patched_source, number_of_aliases_added).""" |
| 74 | count = 0 |
| 75 | lines = source.splitlines(keepends=True) |
| 76 | insertions: list[tuple[int, str]] = [] # (line_idx_after, text_to_insert) |
| 77 | |
| 78 | # Pre-compute which parser variables already have --json defined anywhere |
| 79 | # in the function body (regardless of dest). Argparse will reject duplicates. |
| 80 | already_has_json: set[str] = _vars_with_json(source) |
| 81 | |
| 82 | i = 0 |
| 83 | while i < len(lines): |
| 84 | line = lines[i] |
| 85 | |
| 86 | # Does this line start a .add_argument( call? |
| 87 | m = re.match( |
| 88 | r'^(?P<indent>[ \t]*)(?P<var>\w+)\.add_argument\(', line |
| 89 | ) |
| 90 | if m: |
| 91 | indent = m.group("indent") |
| 92 | var = m.group("var") |
| 93 | |
| 94 | # Always collect the full call body (handles single- and multi-line). |
| 95 | call_lines = [line] |
| 96 | depth = line.count("(") - line.count(")") |
| 97 | |
| 98 | j = i + 1 |
| 99 | while depth > 0 and j < len(lines): |
| 100 | call_lines.append(lines[j]) |
| 101 | depth += lines[j].count("(") - lines[j].count(")") |
| 102 | j += 1 |
| 103 | |
| 104 | call_end = j - 1 # last line of the call (0-based) |
| 105 | call_text = "".join(call_lines) |
| 106 | |
| 107 | # Only act when this is a --format/dest="fmt" call and the parser |
| 108 | # variable doesn't already have a --json argument registered. |
| 109 | if ('"--format"' in call_text |
| 110 | and 'dest="fmt"' in call_text |
| 111 | and var not in already_has_json): |
| 112 | alias = ( |
| 113 | f'{indent}{var}.add_argument(\n' |
| 114 | f'{indent} "--json", action="store_const", const="json", dest="fmt",\n' |
| 115 | f'{indent} help="Shorthand for --format json."\n' |
| 116 | f'{indent})\n' |
| 117 | ) |
| 118 | insertions.append((call_end, alias)) |
| 119 | already_has_json.add(var) |
| 120 | count += 1 |
| 121 | |
| 122 | i = j |
| 123 | continue |
| 124 | |
| 125 | i += 1 |
| 126 | |
| 127 | # Apply insertions in reverse order so line numbers stay valid. |
| 128 | for line_idx, text in sorted(insertions, reverse=True): |
| 129 | lines.insert(line_idx + 1, text) |
| 130 | |
| 131 | return "".join(lines), count |
| 132 | |
| 133 | |
| 134 | def main() -> None: |
| 135 | print("🔍 Querying symbol graph for register() functions in cli/commands…") |
| 136 | raw = _muse( |
| 137 | "code", "query", |
| 138 | "kind=function", "name=register", "file~=cli/commands", |
| 139 | "--json", |
| 140 | ) |
| 141 | data = json.loads(raw) |
| 142 | symbols: list[dict[str, object]] = data["results"] |
| 143 | print(f" Found {len(symbols)} register() functions.\n") |
| 144 | |
| 145 | patched = 0 |
| 146 | skipped = 0 |
| 147 | |
| 148 | for sym in sorted(symbols, key=lambda s: str(s["address"])): |
| 149 | address = str(sym["address"]) |
| 150 | |
| 151 | # Get the exact current source of this function. |
| 152 | source = _muse("code", "cat", address) |
| 153 | |
| 154 | # Strip the header comment line muse prepends ("# file::symbol L…"). |
| 155 | source_lines = source.splitlines(keepends=True) |
| 156 | if source_lines and source_lines[0].startswith("#"): |
| 157 | source_lines = source_lines[1:] |
| 158 | source = "".join(source_lines) |
| 159 | |
| 160 | new_source, added = _add_json_aliases(source) |
| 161 | |
| 162 | if added == 0: |
| 163 | skipped += 1 |
| 164 | continue |
| 165 | |
| 166 | # Write new body to a temp file and patch via muse. |
| 167 | with tempfile.NamedTemporaryFile( |
| 168 | mode="w", suffix=".py", delete=False, encoding="utf-8" |
| 169 | ) as tmp: |
| 170 | tmp.write(new_source) |
| 171 | tmp_path = tmp.name |
| 172 | |
| 173 | result = subprocess.run( |
| 174 | ["muse", "code", "patch", address, "--body", tmp_path, "--json"], |
| 175 | capture_output=True, text=True, |
| 176 | ) |
| 177 | pathlib.Path(tmp_path).unlink(missing_ok=True) |
| 178 | |
| 179 | if result.returncode == 0: |
| 180 | patch_data = json.loads(result.stdout) |
| 181 | print( |
| 182 | f" ✅ {address}\n" |
| 183 | f" +{added} alias(es) added " |
| 184 | f"({patch_data.get('lines_replaced')}→{patch_data.get('new_lines')} lines)" |
| 185 | ) |
| 186 | patched += 1 |
| 187 | else: |
| 188 | print(f" ❌ {address}\n {result.stderr.strip()}", file=sys.stderr) |
| 189 | |
| 190 | print(f"\n{patched} function(s) patched, {skipped} already correct.") |
| 191 | |
| 192 | |
| 193 | if __name__ == "__main__": |
| 194 | main() |
File History
1 commit
sha256:1900655993c83c4107067375548a7be823e471d2515830842f1a12cba4bd3cdf
fix: unified object store migration — idempotent writes, JS…
Sonnet 4.6
minor
⚠
28 days ago