This is an automated email from the ASF dual-hosted git repository.
tqchen pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm-ffi.git
The following commit(s) were added to refs/heads/main by this push:
new 62df2f57 [STUBGEN] Add a `--check` mode that reports stale stub files
(#746)
62df2f57 is described below
commit 62df2f579567802cef3cba0126ac2aea299169ef
Author: Linzhang Li <[email protected]>
AuthorDate: Sat Sep 5 12:33:30 2026 -0400
[STUBGEN] Add a `--check` mode that reports stale stub files (#746)
## Summary
`tvm-ffi-stubgen` updates generated blocks in place, but always exited
0, so CI could not tell whether a tree was up to date. `--check` writes
nothing, prints `[Stale] <file>` for every file whose blocks would
change, and exits 1 if there is any; a file that fails to process exits
2 (in every mode; it used to exit 0). This is the CI half of the
in-place / `--check` pair, like `ruff format --check`.
`--check` is `--dry-run` plus the report: it sets `dry_run` at parse
time, and `_stage_3` returns the changed flag `FileInfo.update` already
computes. It refuses `--init-*`, since `--init` scaffolds files on disk
regardless of `dry_run`.
## Changes
- `stub/cli.py`: `--check` flag, `_stage_3` returns changed, `__main__`
returns `2 if failed else 1 if stale else 0`.
- `stub/utils.py`: `Options.check`.
- `docs/packaging/stubgen.rst`: the option and its exit codes.
- `tests/python/test_stubgen_rust.py`: stale round trip, failure exits
2, `--check` + `--init-*` rejected.
## Testing
`test_stubgen_rust.py` + `test_stubgen.py`: 102 passed; pre-commit lint
clean.
Signed-off-by: yuchuan <[email protected]>
---
docs/packaging/stubgen.rst | 5 +++
python/tvm_ffi/stub/cli.py | 33 +++++++++++++---
python/tvm_ffi/stub/utils.py | 3 +-
tests/python/test_stubgen_rust.py | 82 +++++++++++++++++++++++++++++++++++++++
4 files changed, 116 insertions(+), 7 deletions(-)
diff --git a/docs/packaging/stubgen.rst b/docs/packaging/stubgen.rst
index 1e4ff1ed..3b446a9d 100644
--- a/docs/packaging/stubgen.rst
+++ b/docs/packaging/stubgen.rst
@@ -245,6 +245,11 @@ All three are required together. When omitted, the tool
operates in directive-on
``--dry-run``
Preview changes without writing to files.
+``--check``
+ Write nothing; print ``[Stale] <file>`` for every file whose stub blocks
are out of
+ date and exit with status 1 if there is any (2 if a file failed to
process), so CI
+ can reject a stale tree. Cannot be combined with ``--init-*``.
+
``--imports``
Additional Python modules to import before generation (semicolon-separated).
diff --git a/python/tvm_ffi/stub/cli.py b/python/tvm_ffi/stub/cli.py
index a5e0ba39..a1afe73f 100644
--- a/python/tvm_ffi/stub/cli.py
+++ b/python/tvm_ffi/stub/cli.py
@@ -48,6 +48,9 @@ def __main__() -> int:
This generates in-place type stubs inside special ``tvm-ffi-stubgen``
blocks
in the given files or directories. See the module docstring for an
overview and examples of the block syntax.
+
+ Returns 0, or 2 when a file failed to process, or 1 when ``--check`` found
+ a file whose stub blocks are out of date.
"""
opt = _parse_args()
generator = get_generator(opt.target)
@@ -67,10 +70,13 @@ def __main__() -> int:
# - defined global functions: `tvm-ffi-stubgen(begin): global/...`
# - defined object types: `tvm-ffi-stubgen(begin): object/...`
ty_map: dict[str, str] = generator.default_ty_map()
+ failed = 0
+ stale = 0
for file in files:
try:
_stage_1(file, ty_map)
except Exception:
+ failed += 1
print(
f'{C.TERM_RED}[Failed] File "{file.path}":
{traceback.format_exc()}{C.TERM_RESET}'
)
@@ -95,7 +101,7 @@ def __main__() -> int:
if opt.verbose:
print(f"{C.TERM_CYAN}[File] {file.path}{C.TERM_RESET}")
try:
- _stage_3(
+ changed = _stage_3(
file,
opt,
ty_map,
@@ -103,9 +109,14 @@ def __main__() -> int:
generator=generator,
)
except Exception:
+ failed += 1
print(
f'{C.TERM_RED}[Failed] File "{file.path}":
{traceback.format_exc()}{C.TERM_RESET}'
)
+ continue
+ if changed and opt.check:
+ stale += 1
+ print(f"{C.TERM_YELLOW}[Stale] {file.path}{C.TERM_RESET}")
# Stage 4. Let the generator stitch the generated tree together (runs
after the
# files are fully written, so language-specific wiring isn't clobbered).
@@ -122,7 +133,7 @@ def __main__() -> int:
}
write_coverage_report(Path(opt.coverage_out), classify(infos))
del dlls
- return 0
+ return 2 if failed else 1 if stale else 0
def _stage_1(
@@ -227,7 +238,8 @@ def _stage_3( # noqa: PLR0912
ty_map: dict[str, str],
global_funcs: dict[str, list[FuncInfo]],
generator: Generator,
-) -> None:
+) -> bool:
+ """Process one file's blocks; return whether its content is (or would be)
changed."""
defined_funcs: set[str] = set()
defined_types: set[str] = set()
imports = generator.new_imports()
@@ -273,7 +285,7 @@ def _stage_3( # noqa: PLR0912
if code.kind == "export":
generator.generate_export_block(code)
# Finalize: write back to file
- file.update(verbose=opt.verbose, dry_run=opt.dry_run)
+ return file.update(verbose=opt.verbose, dry_run=opt.dry_run)
def _parse_args() -> Options:
@@ -390,6 +402,14 @@ def _parse_args() -> Options:
"without modifying any files."
),
)
+ parser.add_argument(
+ "--check",
+ action="store_true",
+ help=(
+ "Don't write changes; report every file whose stub blocks are out
of date "
+ "and exit with status 1 if there is any. Cannot be combined with
--init-*."
+ ),
+ )
parser.add_argument(
"--coverage-out",
type=str,
@@ -413,6 +433,8 @@ def _parse_args() -> Options:
shared_target=args.init_lib,
prefix=args.init_prefix,
)
+ if args.check and init_cfg is not None:
+ parser.error("--check cannot be combined with --init-* flags")
if not args.files and args.coverage_out is None:
parser.print_help()
@@ -425,7 +447,8 @@ def _parse_args() -> Options:
indent=args.indent,
files=args.files,
verbose=args.verbose,
- dry_run=args.dry_run,
+ dry_run=args.dry_run or args.check,
+ check=args.check,
target=args.target,
coverage_out=args.coverage_out,
)
diff --git a/python/tvm_ffi/stub/utils.py b/python/tvm_ffi/stub/utils.py
index 7f5ea625..ae22c275 100644
--- a/python/tvm_ffi/stub/utils.py
+++ b/python/tvm_ffi/stub/utils.py
@@ -83,10 +83,9 @@ class Options:
files: list[str] = dataclasses.field(default_factory=list)
verbose: bool = False
dry_run: bool = False
+ check: bool = False
target: str = "python"
- """Code generator target to use."""
coverage_out: str | None = None
- """Path of the native-layout coverage report (JSON) to write, if
requested."""
@dataclasses.dataclass(init=False)
diff --git a/tests/python/test_stubgen_rust.py
b/tests/python/test_stubgen_rust.py
index f2532cb5..438a1131 100644
--- a/tests/python/test_stubgen_rust.py
+++ b/tests/python/test_stubgen_rust.py
@@ -1145,3 +1145,85 @@ def test_cli_init_generates_a_module_tree(tmp_path:
Path, monkeypatch: pytest.Mo
# Running again over the generated tree is a no-op.
assert stub_cli.__main__() == 0
assert (tmp_path / "testing" / "mod.rs").read_text(encoding="utf-8") ==
text
+
+
+def test_cli_check_reports_stale_files(
+ tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys:
pytest.CaptureFixture[str]
+) -> None:
+ monkeypatch.setattr(
+ "sys.argv",
+ [
+ "tvm-ffi-stubgen",
+ "--target",
+ "rust",
+ "--init-pypkg",
+ "demo",
+ "--init-lib",
+ "demo_shared",
+ "--init-prefix",
+ "testing.",
+ str(tmp_path),
+ ],
+ )
+ assert stub_cli.__main__() == 0
+ mod_rs = tmp_path / "testing" / "mod.rs"
+ fresh = mod_rs.read_text(encoding="utf-8")
+ check = ["tvm-ffi-stubgen", "--target", "rust", "--check", str(tmp_path)]
+ # An up-to-date tree passes the check.
+ monkeypatch.setattr("sys.argv", check)
+ assert stub_cli.__main__() == 0
+ assert "[Stale]" not in capsys.readouterr().out
+ # A stale block fails it, is named, and is left untouched.
+ stale = fresh.replace(" pub v_str: String,\n", " pub v_str: String,
// stale\n", 1)
+ assert stale != fresh
+ mod_rs.write_text(stale, encoding="utf-8")
+ assert stub_cli.__main__() == 1
+ assert f"[Stale] {mod_rs}" in capsys.readouterr().out
+ assert mod_rs.read_text(encoding="utf-8") == stale
+ # Running in place repairs it; the check passes again.
+ monkeypatch.setattr("sys.argv", ["tvm-ffi-stubgen", "--target", "rust",
str(tmp_path)])
+ assert stub_cli.__main__() == 0
+ assert mod_rs.read_text(encoding="utf-8") == fresh
+ monkeypatch.setattr("sys.argv", check)
+ assert stub_cli.__main__() == 0
+
+
+def test_cli_failed_file_exits_2(tmp_path: Path, monkeypatch:
pytest.MonkeyPatch) -> None:
+ (tmp_path / "mod.rs").write_text(
+ "\n".join(
+ [
+ f"{C.RUST_SYNTAX.directive('bogus')} testing.TestCxxClassBase",
+ f"{C.RUST_SYNTAX.begin} object/testing.TestCxxClassBase",
+ C.RUST_SYNTAX.end,
+ "",
+ ]
+ ),
+ encoding="utf-8",
+ )
+ for extra in ([], ["--check"]):
+ monkeypatch.setattr(
+ "sys.argv", ["tvm-ffi-stubgen", "--target", "rust", *extra,
str(tmp_path)]
+ )
+ assert stub_cli.__main__() == 2
+
+
+def test_cli_check_rejects_init_flags(tmp_path: Path, monkeypatch:
pytest.MonkeyPatch) -> None:
+ monkeypatch.setattr(
+ "sys.argv",
+ [
+ "tvm-ffi-stubgen",
+ "--target",
+ "rust",
+ "--check",
+ "--init-pypkg",
+ "demo",
+ "--init-lib",
+ "demo_shared",
+ "--init-prefix",
+ "testing.",
+ str(tmp_path),
+ ],
+ )
+ with pytest.raises(SystemExit) as excinfo:
+ stub_cli.__main__()
+ assert excinfo.value.code == 2
