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 13d1ea06 [STUBGEN] Handle object blocks of types that are no longer
registered (#851)
13d1ea06 is described below
commit 13d1ea069c6581ff113d211e5f57dc0a071f2434
Author: Linzhang Li <[email protected]>
AuthorDate: Thu Oct 15 16:39:45 2026 -0400
[STUBGEN] Handle object blocks of types that are no longer registered (#851)
An `object/<key>` block whose type was removed upstream made
`tvm-ffi-stubgen` fail the file with `ValueError: Cannot find type key:
...`. It reported only the first such block per file, and the rest of
the file was not processed.
- In a `prefix` file, a block directly under the prefix whose type is no
longer registered is removed, with a `[Removed]` notice naming the block
and its line. Under `--check` the file is reported as stale.
- Any other such block is left as is. The rest of the file is still
processed and written, then one error lists every such block of the file
with its line, and the run exits with status 2.
Testing:
- `tests/python/test_stubgen_rust.py`: a `prefix` file drops a block of
an unregistered type and passes `--check` afterwards; a file with two
unregistered explicit blocks reports both, by line, and still writes its
other block.
- On TVM's Rust bindings, with stale blocks for
`tirx.transform.StmtSimplifyConfig` and `tirx.tensor.TcMmaAttrs` in
their prefix files and two `tirx.tensor.*` blocks in `tirx.rs`: the
prefix files are restored byte for byte, and both explicit blocks are
listed in one error.
Signed-off-by: yuchuan <[email protected]>
---
docs/packaging/stubgen.rst | 3 ++-
examples/rust_stubgen/README.md | 8 ++++--
python/tvm_ffi/stub/cli.py | 46 +++++++++++++++++++++++++++++---
tests/python/test_stubgen_rust.py | 56 +++++++++++++++++++++++++++++++++++++++
4 files changed, 106 insertions(+), 7 deletions(-)
diff --git a/docs/packaging/stubgen.rst b/docs/packaging/stubgen.rst
index 7b299d46..d2de250a 100644
--- a/docs/packaging/stubgen.rst
+++ b/docs/packaging/stubgen.rst
@@ -391,7 +391,8 @@ When you run the tool, it:
Rust target only. Adds an ``object/<type_key>`` block to the file for every
object
registered directly under the prefix that no processed file defines yet;
``skip``
leaves one out, and a key the ``tvm_ffi`` crate binds itself (every
``ffi.*`` key)
- is never added. Code outside the blocks is preserved.
+ is never added. A block under the prefix whose type is no longer registered
is
+ removed. Code outside the blocks is preserved.
.. code-block:: rust
diff --git a/examples/rust_stubgen/README.md b/examples/rust_stubgen/README.md
index b855bbe3..599de016 100644
--- a/examples/rust_stubgen/README.md
+++ b/examples/rust_stubgen/README.md
@@ -75,7 +75,8 @@ directly under it gets an `object/<type_key>` block in this
file on the next
run, and a `skip` line (`// tvm-ffi-stubgen(skip): rust_stubgen.Internal`)
leaves one out. Add a second object to `int_pair.cc` and rebuild, and its block
appears after `IntPair`; delete a block and it comes back, so dropping a
binding
-means writing `skip`. Code outside the blocks is never touched.
+means writing `skip`. Remove the object from the C++ side and its block is
+removed on the next run. Code outside the blocks is never touched.
Nothing requires a `generated/` directory: the blocks can live in any `.rs`
file the command line points at. This example keeps them under
@@ -146,4 +147,7 @@ importing it.
```
4. In CI, run the same command with `--check`. It writes nothing and exits with
- status 1 when a block is out of date (2 when a file fails to process).
+ status 1 when a block is out of date (2 when a file fails to process). A
+ block whose type is no longer registered, outside a `prefix` file's own
+ namespace, fails its file; every such block of the file is listed with its
+ line.
diff --git a/python/tvm_ffi/stub/cli.py b/python/tvm_ffi/stub/cli.py
index af8f8c24..24e1d5b5 100644
--- a/python/tvm_ffi/stub/cli.py
+++ b/python/tvm_ffi/stub/cli.py
@@ -26,6 +26,8 @@ import traceback
from pathlib import Path
from typing import TYPE_CHECKING
+from tvm_ffi.core import _object_type_key_to_index
+
from . import consts as C
from .file_utils import CodeBlock, FileInfo, collect_files, syntax_for
from .generator import generator_names, get_generator
@@ -170,8 +172,9 @@ def _roll_out_prefixes(files: list[FileInfo], generator:
Generator) -> int:
"""Append an ``object/<key>`` block for each registered object under a
file's ``prefix``.
Keys with a block in any file of the run, named by ``skip``, or bound by
the target's
- runtime are left alone; an ``import-section`` is added when the file has
none. Returns
- the number of bad files.
+ runtime are left alone; an ``import-section`` is added when the file has
none. A block
+ under the prefix whose type is no longer registered is removed. Returns
the number of
+ bad files.
"""
defined = {code.param for file in files for code in file.code_blocks if
code.kind == "object"}
registry = collect_type_keys()
@@ -193,6 +196,7 @@ def _roll_out_prefixes(files: list[FileInfo], generator:
Generator) -> int:
failed += 1
print(f'{C.TERM_RED}[Failed] File "{file.path}":
{error}{C.TERM_RESET}')
continue
+ _remove_unregistered_blocks(file, prefix, registry.get(prefix, []))
if prefix not in registry:
print(
f"{C.TERM_YELLOW}[Skipped] No registered object under prefix
`{prefix}`{C.TERM_RESET}"
@@ -217,6 +221,27 @@ def _roll_out_prefixes(files: list[FileInfo], generator:
Generator) -> int:
return failed
+def _remove_unregistered_blocks(file: FileInfo, prefix: str, registered:
list[str]) -> None:
+ """Drop each ``object/`` block directly under ``prefix`` whose type is not
``registered``."""
+ kept: list[CodeBlock] = []
+ for code in file.code_blocks:
+ if (
+ code.kind == "object"
+ and isinstance(code.param, str)
+ and code.param.rpartition(".")[0] == prefix
+ and code.param not in registered
+ ):
+ print(
+ f'{C.TERM_YELLOW}[Removed] File "{file.path}":
`object/{code.param}` '
+ f"(line {code.lineno_start}): the type is no longer
registered{C.TERM_RESET}"
+ )
+ if kept and kept[-1].kind is None and kept[-1].lines == [""]:
+ kept.pop() # the blank line the roll-out put before it
+ continue
+ kept.append(code)
+ file.code_blocks[:] = kept
+
+
def _new_blocks(syntax: C.MarkerSyntax, lineno: int, stub: str) ->
list[CodeBlock]:
"""Return a blank line and an empty ``begin``/``end`` block for ``stub``,
to insert into a file."""
begin = f"{syntax.begin} {stub}"
@@ -323,8 +348,13 @@ def _stage_3( # noqa: PLR0912
generator: Generator,
declared: Container[str] = frozenset(),
) -> bool:
- """Process one file's blocks; return whether its content is (or would be)
changed."""
+ """Process one file's blocks; return whether its content is (or would be)
changed.
+
+ An ``object/`` block whose type is no longer registered is left as is;
once the rest
+ of the file is processed, they are reported together as an error.
+ """
defined_funcs: set[str] = set()
+ unregistered: list[str] = []
defined_types: set[str] = set()
imports = generator.new_imports()
# Stage 1. Hand the one-line directives the pipeline does not consume
itself to the generator.
@@ -349,6 +379,9 @@ def _stage_3( # noqa: PLR0912
if code.kind == "object":
type_key = code.param
assert isinstance(type_key, str)
+ if _object_type_key_to_index(type_key) is None:
+ unregistered.append(f"`object/{type_key}` (line
{code.lineno_start})")
+ continue
obj_info = object_info_from_type_key(type_key)
type_key = ty_map.get(type_key, type_key)
defined_types.add(generator.canonical_type_name(type_key))
@@ -369,7 +402,12 @@ def _stage_3( # noqa: PLR0912
if code.kind == "export":
generator.generate_export_block(code)
# Finalize: write back to file
- return file.update(verbose=opt.verbose, dry_run=opt.dry_run)
+ changed = file.update(verbose=opt.verbose, dry_run=opt.dry_run)
+ if unregistered:
+ raise ValueError(
+ f"{', '.join(unregistered)}: the type is no longer registered;
delete the block"
+ )
+ return changed
def _parse_args() -> Options:
diff --git a/tests/python/test_stubgen_rust.py
b/tests/python/test_stubgen_rust.py
index 17689c12..8f5ff12a 100644
--- a/tests/python/test_stubgen_rust.py
+++ b/tests/python/test_stubgen_rust.py
@@ -1646,3 +1646,59 @@ def test_prefix_survives_init(tmp_path: Path,
monkeypatch: pytest.MonkeyPatch) -
"sys.argv", ["tvm-ffi-stubgen", "--target", "rust", "--check",
str(tmp_path)]
)
assert stub_cli.__main__() == 0
+
+
+# ---------------------------------------------------------------------------
+# Blocks whose type is no longer registered
+# ---------------------------------------------------------------------------
+
+
+def test_prefix_removes_blocks_of_unregistered_types(
+ tmp_path: Path, monkeypatch: pytest.MonkeyPatch, capsys:
pytest.CaptureFixture[str]
+) -> None:
+ src = tmp_path / "testing.rs"
+ gone = f"{C.RUST_SYNTAX.begin} object/testing.NoLongerRegistered"
+ src.write_text(
+ "\n".join(
+ [f"{C.RUST_SYNTAX.directive('prefix')} testing", "", gone,
C.RUST_SYNTAX.end, ""]
+ ),
+ encoding="utf-8",
+ )
+ monkeypatch.setattr("sys.argv", ["tvm-ffi-stubgen", "--target", "rust",
str(tmp_path)])
+ assert stub_cli.__main__() == 0
+ out = capsys.readouterr().out
+ assert "[Removed]" in out and "`object/testing.NoLongerRegistered` (line
3)" in out
+ text = src.read_text(encoding="utf-8")
+ assert "NoLongerRegistered" not in text
+ assert "pub struct TestCxxClassBaseObj {" in text # the rest of the
namespace is rolled out
+ monkeypatch.setattr(
+ "sys.argv", ["tvm-ffi-stubgen", "--target", "rust", "--check",
str(tmp_path)]
+ )
+ assert stub_cli.__main__() == 0
+
+
+def test_explicit_blocks_of_unregistered_types_are_reported_together(tmp_path:
Path) -> None:
+ src = tmp_path / "mod.rs"
+ blocks = [
+ f"{C.RUST_SYNTAX.begin} import-section",
+ C.RUST_SYNTAX.end,
+ f"{C.RUST_SYNTAX.begin} object/demo.Gone",
+ C.RUST_SYNTAX.end,
+ f"{C.RUST_SYNTAX.begin} object/testing.TestCxxClassBase",
+ C.RUST_SYNTAX.end,
+ f"{C.RUST_SYNTAX.begin} object/demo.AlsoGone",
+ C.RUST_SYNTAX.end,
+ "",
+ ]
+ src.write_text("\n".join(blocks), encoding="utf-8")
+ info = FileInfo.from_file(src)
+ assert info is not None
+ expected = (
+ r"`object/demo\.Gone` \(line 3\), `object/demo\.AlsoGone` \(line 7\): "
+ r"the type is no longer registered"
+ )
+ with pytest.raises(ValueError, match=expected):
+ _stage_3(info, Options(), RUST.default_ty_map(), {}, RUST,
ALL_DECLARED)
+ text = src.read_text(encoding="utf-8")
+ assert "pub struct TestCxxClassBaseObj {" in text # the other blocks are
still written
+ assert f"{C.RUST_SYNTAX.begin} object/demo.Gone\n{C.RUST_SYNTAX.end}" in
text
