mirror of
https://github.com/harvard-edge/cs249r_book.git
synced 2026-07-16 14:42:29 -05:00
399 lines
16 KiB
Python
399 lines
16 KiB
Python
"""Tests for the paired multiplier codemod (provable lane + queue lane)."""
|
||
import sys
|
||
from pathlib import Path
|
||
|
||
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "tools/audit/fmt"))
|
||
import ast # noqa: E402
|
||
|
||
from codemod_fmt import ( # noqa: E402
|
||
_rewrite_call_to_multiple, _patch_prose, scan_file,
|
||
_rewrite_percent, scan_percent,
|
||
_rewrite_scale, scan_scale,
|
||
_rewrite_scale_style, scan_scale_style,
|
||
_rewrite_unit, scan_unit,
|
||
_rewrite_time, scan_time,
|
||
_rewrite_rate, scan_rate,
|
||
)
|
||
|
||
|
||
def _percent(expr: str):
|
||
"""Parse a single fmt(...) expression and run the percent rewriter on it."""
|
||
call = ast.parse(expr, mode="eval").body
|
||
return _rewrite_percent(call, expr)
|
||
|
||
|
||
def _scale(expr: str):
|
||
call = ast.parse(expr, mode="eval").body
|
||
return _rewrite_scale(call, expr)
|
||
|
||
|
||
def _scale_style(expr: str):
|
||
call = ast.parse(expr, mode="eval").body
|
||
return _rewrite_scale_style(call, expr)
|
||
|
||
|
||
def _unit(expr: str):
|
||
call = ast.parse(expr, mode="eval").body
|
||
return _rewrite_unit(call, expr)
|
||
|
||
|
||
def _time(expr: str):
|
||
call = ast.parse(expr, mode="eval").body
|
||
return _rewrite_time(call, expr)
|
||
|
||
|
||
def _rate(expr: str):
|
||
call = ast.parse(expr, mode="eval").body
|
||
return _rewrite_rate(call, expr)
|
||
|
||
|
||
# --- the provable cell rewrite -------------------------------------------------
|
||
|
||
def test_fmt_multiplier_rewrite_drops_suffix_keeps_precision():
|
||
out = _rewrite_call_to_multiple("fmt(adam_mult, precision=0, commas=False, suffix='×')")
|
||
assert out == "fmt_multiple(adam_mult, precision=0, commas=False)"
|
||
|
||
|
||
def test_fmt_int_is_not_auto_rewritten():
|
||
# fmt_int rounds; not provably equal to fmt_multiple -> declined here
|
||
assert _rewrite_call_to_multiple("fmt_int(round(x), commas=False, suffix='×')") is None
|
||
|
||
|
||
def test_multiline_call_declined():
|
||
assert _rewrite_call_to_multiple("fmt(x,\n suffix='×')") is None
|
||
|
||
|
||
def test_commas_true_injected_when_omitted():
|
||
# fmt defaults commas=True; fmt_multiple defaults commas=False -> must pin
|
||
assert _rewrite_call_to_multiple("fmt(x, suffix='×')") == "fmt_multiple(x, commas=True)"
|
||
|
||
|
||
def test_explicit_commas_false_preserved():
|
||
out = _rewrite_call_to_multiple("fmt(x, commas=False, suffix='×')")
|
||
assert out == "fmt_multiple(x, commas=False)"
|
||
|
||
|
||
def test_prefix_call_declined():
|
||
# fmt_multiple has no prefix= -> not a clean rewrite
|
||
assert _rewrite_call_to_multiple("fmt(x, prefix='~', suffix='×')") is None
|
||
|
||
|
||
# --- the prose patch (position-based, handles repeats) -------------------------
|
||
|
||
def test_prose_patch_single_ref():
|
||
text = "speedup is `{python} C.s_str` over baseline\n"
|
||
new, patches = _patch_prose(text, {"C.s_str"})
|
||
assert "`{python} C.s_str`$\\times$ over baseline" in new
|
||
assert len(patches) == 1
|
||
|
||
|
||
def test_prose_patch_repeated_ref_on_one_line_each_once():
|
||
# the bug that produced '16××' / '1.75×××' must not recur
|
||
text = "| h | `{python} C.s_str` | `{python} C.s_str` |\n"
|
||
new, patches = _patch_prose(text, {"C.s_str"})
|
||
assert new.count("$\\times$") == 2
|
||
assert "$\\times$$\\times$" not in new
|
||
assert len(patches) == 2
|
||
|
||
|
||
def test_prose_patch_skips_when_glyph_already_present():
|
||
text = "ratio `{python} C.s_str`$\\times$ already\n"
|
||
new, patches = _patch_prose(text, {"C.s_str"})
|
||
assert patches == []
|
||
assert new.count("$\\times$") == 1
|
||
|
||
|
||
def test_prose_patch_ignores_cell_bodies():
|
||
text = "```{python}\ns_str = fmt_multiple(6)\n```\nvalue `{python} C.s_str` here\n"
|
||
new, patches = _patch_prose(text, {"C.s_str"})
|
||
# only the prose ref is patched, not the assignment in the cell
|
||
assert "fmt_multiple(6)$\\times$" not in new
|
||
assert len(patches) == 1
|
||
|
||
|
||
def test_prose_patch_ignores_unrelated_refs():
|
||
text = "cost `{python} C.cost_str` total\n"
|
||
new, patches = _patch_prose(text, {"C.s_str"})
|
||
assert patches == []
|
||
assert new == text
|
||
|
||
|
||
# --- queue lane ----------------------------------------------------------------
|
||
|
||
CELL = """```{{python}}
|
||
#| echo: false
|
||
class C:
|
||
{assigns}
|
||
```
|
||
|
||
{prose}
|
||
"""
|
||
|
||
|
||
def test_percent_and_scale_go_to_queue_not_rewritten(tmp_path):
|
||
p = tmp_path / "c.qmd"
|
||
p.write_text(CELL.format(
|
||
assigns=("pct_str = fmt(savings, suffix='%')\n"
|
||
" big_str = fmt(count / 1e6, suffix='M')\n"
|
||
" sp_str = fmt(speedup, suffix='×')"),
|
||
prose="x `{python} C.sp_str` y"), encoding="utf-8")
|
||
edits, mult_vars, queue = scan_file(p)
|
||
kinds = {q.kind for q in queue}
|
||
assert "percent" in kinds and "scale" in kinds
|
||
# only the multiplier is auto-rewritten
|
||
assert [e.var for e in edits] == ["sp_str"]
|
||
assert mult_vars == {"C.sp_str"}
|
||
|
||
|
||
def test_cross_class_same_name_only_rewritten_one_is_patched(tmp_path):
|
||
# the data_engineering bug: two classes export 'ratio_str'; only Auto's is
|
||
# rewritten (fmt …×), Kept's uses fmt_int (queued). Prose must patch ONLY
|
||
# Auto.ratio_str, never Kept.ratio_str (which still carries × in its string).
|
||
body = (
|
||
"```{python}\n#| echo: false\n"
|
||
"class Auto:\n ratio_str = fmt(a, precision=1, commas=False, suffix='×')\n"
|
||
"class Kept:\n ratio_str = fmt_int(round(b), commas=False, suffix='×')\n"
|
||
"```\n\n"
|
||
"auto `{python} Auto.ratio_str` and kept `{python} Kept.ratio_str` end\n"
|
||
)
|
||
p = tmp_path / "c.qmd"
|
||
p.write_text(body, encoding="utf-8")
|
||
edits, mult_vars, queue = scan_file(p)
|
||
assert mult_vars == {"Auto.ratio_str"}
|
||
new, patches = _patch_prose(body, mult_vars)
|
||
assert "`{python} Auto.ratio_str`$\\times$" in new
|
||
assert "`{python} Kept.ratio_str`$\\times$" not in new
|
||
assert [r for _, r in patches] == ["Auto.ratio_str"]
|
||
|
||
|
||
def test_literal_x_and_space_glyph_go_to_queue_not_auto(tmp_path):
|
||
p = tmp_path / "c.qmd"
|
||
p.write_text(CELL.format(
|
||
assigns=("xx_str = fmt(ratio, suffix='x')\n"
|
||
" sp_str = fmt(speedup, suffix=' ×')"),
|
||
prose="x"), encoding="utf-8")
|
||
edits, mult_vars, queue = scan_file(p)
|
||
assert edits == [] # neither auto-rewritten
|
||
assert all(q.kind == "multiplier-variant" for q in queue)
|
||
assert len(queue) == 2
|
||
|
||
|
||
# --- the byte-identical percent lane -------------------------------------------
|
||
|
||
def test_percent_symbol_standalone_divides_by_100():
|
||
out = _percent("fmt(savings_pct, precision=1, commas=False, suffix='%')")
|
||
assert out == "fmt_percent(savings_pct/100, precision=1, commas=False, style='symbol')"
|
||
|
||
|
||
def test_percent_prose_word_uses_prose_style():
|
||
out = _percent("fmt(share, precision=0, suffix=' percent')")
|
||
assert out == "fmt_percent(share/100, precision=0, style='prose')"
|
||
|
||
|
||
def test_percent_strips_ratio_times_100():
|
||
# the clean win: fmt(ratio*100, '%') -> fmt_percent(ratio) (no *100/100)
|
||
assert _percent("fmt(mfu * 100, precision=1, suffix='%')") == \
|
||
"fmt_percent(mfu, precision=1, style='symbol')"
|
||
assert _percent("fmt(100 * util, suffix='%')") == "fmt_percent(util, style='symbol')"
|
||
|
||
|
||
def test_percent_parenthesises_binop_before_dividing():
|
||
# a bare expression must be parenthesised so /100 binds to the whole value
|
||
assert _percent("fmt(a + b, suffix='%')") == "fmt_percent((a + b)/100, style='symbol')"
|
||
|
||
|
||
def test_percent_declines_prefix_and_nonexact_suffix():
|
||
assert _percent("fmt(x, prefix='~', suffix='%')") is None
|
||
assert _percent("fmt(x, suffix=' %')") is None # spacing variant -> queue
|
||
assert _percent("fmt(x, suffix='percent')") is None # no leading space -> queue
|
||
|
||
|
||
def test_scan_percent_migrates_fmt_int_exact_queues_variants(tmp_path):
|
||
p = tmp_path / "c.qmd"
|
||
p.write_text(CELL.format(
|
||
assigns=("ok_str = fmt(acc, suffix='%')\n"
|
||
" int_str = fmt_int(rate, suffix='%')\n"
|
||
" sp_str = fmt(x, suffix=' %')"),
|
||
prose="x"), encoding="utf-8")
|
||
edits, queue = scan_percent(p)
|
||
# exact fmt('%') and fmt_int('%') both migrate (the latter via round()/100)
|
||
assert sorted(e.var for e in edits) == ["int_str", "ok_str"]
|
||
int_edit = next(e for e in edits if e.var == "int_str")
|
||
assert int_edit.new == "fmt_percent(round(rate)/100, precision=0, commas=True, style='symbol')"
|
||
# the spacing variant ' %' still goes to the adjudication queue
|
||
assert [q.var for q in queue] == ["sp_str"]
|
||
assert queue[0].kind == "percent"
|
||
|
||
|
||
def test_scale_strips_named_divisor_matching_glyph():
|
||
assert _scale("fmt(x / MILLION, suffix='M')") == \
|
||
"fmt_count(x, scale='M', precision=1)"
|
||
assert _scale("fmt(toks / c.TRILLION, precision=1, commas=False, suffix='T')") == \
|
||
"fmt_count(toks, scale='T', precision=1, commas=False)"
|
||
|
||
|
||
def test_scale_strips_numeric_literal_divisor():
|
||
assert _scale("fmt(n / 1e9, precision=2, suffix='B')") == \
|
||
"fmt_count(n, scale='B', precision=2)"
|
||
|
||
|
||
def test_scale_declines_when_divisor_does_not_match_glyph():
|
||
# divisor magnitude (1e3) != glyph 'M' (1e6): a latent bug, leave it
|
||
assert _scale("fmt(x / THOUSAND, suffix='M')") is None
|
||
|
||
|
||
def test_scale_declines_prescaled_and_prefix_and_lowercase():
|
||
assert _scale("fmt(params_b, suffix='B')") is None # no division
|
||
assert _scale("fmt(x / THOUSAND, suffix='k')") is None # lowercase glyph
|
||
assert _scale("fmt(x / MILLION, prefix='~', suffix='M')") is None
|
||
|
||
|
||
def test_scan_scale_queues_prescaled_migrates_division(tmp_path):
|
||
p = tmp_path / "c.qmd"
|
||
p.write_text(CELL.format(
|
||
assigns=("a_str = fmt(toks / MILLION, suffix='M')\n"
|
||
" b_str = fmt(params_b, suffix='B')\n"
|
||
" c_str = fmt_int(n / THOUSAND, suffix='K')"),
|
||
prose="x"), encoding="utf-8")
|
||
edits, queue = scan_scale(p)
|
||
assert [e.var for e in edits] == ["a_str"] # only the clean division
|
||
assert sorted(q.var for q in queue) == ["b_str", "c_str"]
|
||
assert all(q.kind == "scale" for q in queue)
|
||
|
||
|
||
def test_scale_style_prescaled_value_reconstructs_raw_count():
|
||
assert _scale_style("fmt(params_b, precision=0, commas=False, suffix=' B')") == \
|
||
"fmt_count(params_b * BILLION, scale='B', precision=0, commas=False)"
|
||
|
||
|
||
def test_scale_style_lowercase_k_division_becomes_uppercase_scale():
|
||
assert _scale_style("fmt(n / THOUSAND, precision=0, commas=False, suffix='k')") == \
|
||
"fmt_count(n, scale='K', precision=0, commas=False)"
|
||
|
||
|
||
def test_scale_style_param_quantity_uses_raw_param_count():
|
||
assert _scale_style("fmt(model.parameters.m_as(Bparam), precision=0, commas=False, suffix=' B')") == \
|
||
"fmt_count(model.parameters.m_as('param'), scale='B', precision=0, commas=False)"
|
||
|
||
|
||
def test_scale_style_fmt_int_round_division():
|
||
assert _scale_style("fmt_int(round(n_params / BILLION), commas=False, suffix='B')") == \
|
||
"fmt_count(round(n_params / BILLION) * BILLION, scale='B', precision=0, commas=False)"
|
||
|
||
|
||
def test_scan_scale_style_qualifies_class_exports(tmp_path):
|
||
p = tmp_path / "c.qmd"
|
||
p.write_text(CELL.format(
|
||
assigns=("a_str = fmt(params_b, precision=0, commas=False, suffix=' B')\n"
|
||
" b_str = fmt(n / THOUSAND, precision=0, commas=False, suffix='k')"),
|
||
prose="x"), encoding="utf-8")
|
||
edits, qualified, queue = scan_scale_style(p)
|
||
assert sorted(e.var for e in edits) == ["a_str", "b_str"]
|
||
assert qualified == {"C.a_str", "C.b_str"}
|
||
assert queue == []
|
||
|
||
|
||
def test_scan_percent_queues_multiline_calls_not_silently_dropped(tmp_path):
|
||
# a multiline fmt(...) can't be spliced by the single-line applier; it must
|
||
# land in the queue, never be emitted as an edit that won't apply.
|
||
p = tmp_path / "c.qmd"
|
||
p.write_text(
|
||
"```{python}\n#| echo: false\nclass C:\n"
|
||
" total_str = fmt(\n (a + b) * 100,\n precision=0,\n"
|
||
" suffix=' percent',\n )\n```\n\nx `{python} C.total_str` y\n",
|
||
encoding="utf-8")
|
||
edits, queue = scan_percent(p)
|
||
assert edits == []
|
||
assert len(queue) == 1 and queue[0].kind == "percent"
|
||
|
||
|
||
def test_unit_rewrite_quantity_m_as_to_fmt_qty():
|
||
assert _unit("fmt(d_vol.m_as('GB'), precision=0, commas=False, suffix=' GB')") == \
|
||
"fmt_qty(d_vol, GB, precision=0, commas=False)"
|
||
assert _unit("fmt(bw.m_as(TB/second), precision=2, commas=False, suffix=' TB/s')") == \
|
||
"fmt_qty(bw, TB/second, precision=2, commas=False)"
|
||
assert _unit("fmt(energy.m_as(joule), precision=2, commas=False, suffix=' J/token')") == \
|
||
"fmt_qty(energy, joule, precision=2, commas=False, per='token')"
|
||
|
||
|
||
def test_unit_rewrite_singularizes_flop_rate_suffix():
|
||
assert _unit("fmt(gpu.compute.peak_flops.m_as('TFLOPs/s'), precision=0, commas=False, suffix=' TFLOP/s')") == \
|
||
"fmt_qty(gpu.compute.peak_flops, TFLOP/second, precision=0, commas=False)"
|
||
|
||
|
||
def test_unit_rewrite_pins_default_commas():
|
||
assert _unit("fmt(mem.m_as(GB), suffix=' GB')") == "fmt_qty(mem, GB, commas=True)"
|
||
|
||
|
||
def test_unit_rewrite_declines_plain_float_and_noncanonical_words():
|
||
assert _unit("fmt(weights_gb, precision=1, suffix=' GB')") is None
|
||
assert _unit("fmt(cluster_mtbf.m_as(ureg.hour), precision=1, suffix=' hours')") is None
|
||
|
||
|
||
def test_scan_unit_migrates_quantity_backed_and_queues_plain_float(tmp_path):
|
||
p = tmp_path / "c.qmd"
|
||
p.write_text(CELL.format(
|
||
assigns=("ok_str = fmt(mem.m_as(GB), precision=0, commas=False, suffix=' GB')\n"
|
||
" float_str = fmt(mem_gb, precision=0, suffix=' GB')"),
|
||
prose="x"), encoding="utf-8")
|
||
edits, queue = scan_unit(p)
|
||
assert [e.var for e in edits] == ["ok_str"]
|
||
assert edits[0].new == "fmt_qty(mem, GB, precision=0, commas=False)"
|
||
assert [q.var for q in queue] == ["float_str"]
|
||
assert queue[0].kind == "unit"
|
||
|
||
|
||
def test_time_rewrite_uses_full_unit_names_for_symbol_style():
|
||
assert _time("fmt(latency_ms, precision=0, commas=False, suffix=' ms')") == \
|
||
"fmt_time(latency_ms, 'millisecond', precision=0, commas=False)"
|
||
assert _time("fmt(duration_s, precision=1, suffix=' s')") == \
|
||
"fmt_time(duration_s, 'second', precision=1, commas=True)"
|
||
|
||
|
||
def test_time_rewrite_word_style_and_fmt_int_rounding():
|
||
assert _time("fmt(elapsed, precision=0, commas=False, suffix=' seconds')") == \
|
||
"fmt_time(elapsed, 'second', precision=0, commas=False, style='word')"
|
||
assert _time("fmt(latency, precision=0, commas=False, suffix=' milliseconds')") == \
|
||
"fmt_time(latency, 'millisecond', precision=0, commas=False, style='word')"
|
||
assert _time("fmt(us, precision=1, commas=False, suffix=' microseconds')") == \
|
||
"fmt_time(us, 'microsecond', precision=1, commas=False, style='word')"
|
||
assert _time("fmt_int(hours, commas=False, suffix=' hours')") == \
|
||
"fmt_time(round(hours), 'hour', precision=0, commas=False, style='word')"
|
||
|
||
|
||
def test_time_rewrite_preserves_markers():
|
||
assert _time("fmt(duration, precision=0, commas=False, approx=True, suffix=' s')") == \
|
||
"fmt_time(duration, 'second', precision=0, commas=False, approx=True)"
|
||
|
||
|
||
def test_scan_time_migrates_exact_time_suffixes(tmp_path):
|
||
p = tmp_path / "c.qmd"
|
||
p.write_text(CELL.format(
|
||
assigns=("latency_str = fmt(latency, precision=0, commas=False, suffix=' ms')\n"
|
||
" duration_str = fmt_int(duration, commas=True, suffix=' seconds')\n"
|
||
" other_str = fmt(size, precision=0, suffix=' GB')"),
|
||
prose="x"), encoding="utf-8")
|
||
edits, queue = scan_time(p)
|
||
assert [e.var for e in edits] == ["latency_str", "duration_str"]
|
||
assert queue == []
|
||
|
||
|
||
def test_rate_rewrite_service_rate_suffixes():
|
||
assert _rate("fmt(tps, precision=0, commas=False, suffix=' tokens/s')") == \
|
||
"fmt_rate(tps, 'tokens/s', precision=0, commas=False)"
|
||
assert _rate("fmt(fps, precision=0, commas=False, suffix=' FPS')") == \
|
||
"fmt_rate(fps, 'FPS', precision=0, commas=False)"
|
||
assert _rate("fmt_int(imgs, suffix=' img/s')") == \
|
||
"fmt_rate(round(imgs), 'img/s', precision=0)"
|
||
|
||
|
||
def test_scan_rate_migrates_exact_service_rates(tmp_path):
|
||
p = tmp_path / "c.qmd"
|
||
p.write_text(CELL.format(
|
||
assigns=("toks_str = fmt(toks, precision=0, commas=False, suffix=' tokens/s')\n"
|
||
" bw_str = fmt(bw, precision=0, commas=False, suffix=' GB/s')"),
|
||
prose="x"), encoding="utf-8")
|
||
edits, queue = scan_rate(p)
|
||
assert [e.var for e in edits] == ["toks_str"]
|
||
assert queue == []
|