diff --git a/pixi.toml b/pixi.toml index c30864607e6c..8b35e800529f 100644 --- a/pixi.toml +++ b/pixi.toml @@ -354,6 +354,8 @@ py-test = { cmd = "python -m pytest -vv rerun_py/tests/unit", depends-on = [ "py-build", ] } +snapshots = "python scripts/snapshots.py" + [feature.python-docs.tasks] # Build the documentation search index. # See `pixi run search-index --help` for more information. diff --git a/scripts/snapshots.py b/scripts/snapshots.py new file mode 100644 index 000000000000..556454ccd03a --- /dev/null +++ b/scripts/snapshots.py @@ -0,0 +1,128 @@ +""" +View or clean failed kittest snapshot tests. + +Usage: +``` +pixi run snapshot --help +``` +""" + +from __future__ import annotations + +import argparse +from pathlib import Path +from sys import stderr +from typing import Iterator + +import numpy as np +import PIL.Image as Image +import rerun as rr +import rerun.blueprint as rrb + +CRATES_DIR = Path(__file__).parent.parent / "crates" + + +def find_failed_snapshot_tests(package: str | None) -> Iterator[tuple[Path, Path, Path]]: + for diff_path in CRATES_DIR.rglob("**/*.diff.png"): + if package is not None: + if not any(package == str(part) for part in diff_path.parts): + continue + + original_path = diff_path.parent / diff_path.name.replace(".diff.png", ".png") + new_path = diff_path.parent / diff_path.name.replace(".diff.png", ".new.png") + + if original_path.exists() and new_path.exists(): + yield original_path, new_path, diff_path + + +def blueprint(path: Path) -> rrb.Blueprint: + test_name = path.stem + crate_name = path.relative_to(CRATES_DIR).parts[1] + + return rrb.Blueprint( + rrb.Tabs( + rrb.Horizontal( + rrb.Spatial2DView(origin="original", name="Original"), + rrb.Spatial2DView(origin="new", name="New"), + rrb.Spatial2DView(origin="diff", name="Diff"), + name="Side-by-side", + ), + rrb.Tabs( + rrb.Spatial2DView(origin="original", name="Original"), + rrb.Spatial2DView(origin="new", name="New"), + rrb.Spatial2DView(origin="diff", name="Diff"), + ), + rrb.Tabs( + rrb.Vertical( + rrb.Spatial2DView( + contents=["/original", "/new"], + name="Overlay (opacity)", + overrides={ + "/new": [rr.components.Opacity(0.5)], + }, + ), + name='NOTE: Select the "new" entity visualizer and play with the "Opacity" component', + ), + name="Overlay (tab)", + ), + name=f"{crate_name}/{test_name}", + ), + rrb.TimePanel(expanded=False), + ) + + +def log_failed_snapshot_tests(original_path: Path, new_path: Path, diff_path: Path, args: argparse.Namespace) -> None: + recording = rr.new_recording(f"rerun_example_{original_path.stem}") + + with recording: + default_blueprint = blueprint(original_path) + + if args.stdout: + rr.stdout(default_blueprint=default_blueprint) + elif args.serve: + rr.serve(default_blueprint=default_blueprint) + elif args.connect: + rr.connect(args.addr, default_blueprint=default_blueprint) + elif args.save is not None: + rr.save(args.save, default_blueprint=default_blueprint) + elif not args.headless: + rr.spawn(default_blueprint=default_blueprint) + + rr.log("original", rr.Image(np.array(Image.open(original_path))), static=True) + rr.log("new", rr.Image(np.array(Image.open(new_path))), static=True) + rr.log("diff", rr.Image(np.array(Image.open(diff_path))), static=True) + + rr.log( + "doc/tabs", + rr.TextDocument( + "### Click on one of the tabs below to show the Original/New/Diff images.", media_type="text/markdown" + ), + ) + + +def main(): + parser = argparse.ArgumentParser(description="Logs all failed snapshot tests for comparison in rerun") + parser.add_argument("-p", "--package", type=str, help="Only consider the provided package") + parser.add_argument("--clean", action="store_true", help="Clean snapshot files instead of displaying them") + rr.script_add_args(parser) + + args = parser.parse_args() + + none_found = True + for original_path, new_path, diff_path in find_failed_snapshot_tests(args.package): + none_found = False + + if args.clean: + print(f"Removing {new_path}", file=stderr) + new_path.unlink(missing_ok=True) + print(f"Removing {diff_path}", file=stderr) + diff_path.unlink(missing_ok=True) + else: + log_failed_snapshot_tests(original_path, new_path, diff_path, args) + + if none_found: + print("No failed snapshot found", file=stderr) + + +if __name__ == "__main__": + main()