From 1e4b20be54aeb70a22fa0b142a0c9aa131d14dd3 Mon Sep 17 00:00:00 2001 From: yuchuan Date: Sat, 5 Sep 2026 12:09:39 -0400 Subject: [PATCH 1/4] update. Signed-off-by: yuchuan --- docs/packaging/stubgen.rst | 19 ++++++ python/tvm_ffi/stub/cli.py | 62 +++++++++++++++++- python/tvm_ffi/stub/consts.py | 2 +- tests/python/test_stubgen_rust.py | 101 ++++++++++++++++++++++++++++++ 4 files changed, 182 insertions(+), 2 deletions(-) diff --git a/docs/packaging/stubgen.rst b/docs/packaging/stubgen.rst index 3b446a9d0..a88cf29b8 100644 --- a/docs/packaging/stubgen.rst +++ b/docs/packaging/stubgen.rst @@ -387,6 +387,25 @@ When you run the tool, it: # tvm-ffi-stubgen(ty-map): ffi.reflection.AccessStep -> ffi.access_path.AccessStep +``prefix`` - Demand a Namespace + Declares that the file holds every object registered directly under a type-key + prefix (``tirx`` covers ``tirx.Add`` but not ``tirx.transform.X``). Each run appends + an empty ``object/`` block for every such object that no processed file + defines yet, parents first, after the file's last block (an ``import-section`` is + added after the ``prefix`` line when the file has none), and then fills it. Code + outside the blocks is preserved, so a file can start as a one-line skeleton. + + .. code-block:: python + + # tvm-ffi-stubgen(prefix): my_ffi_extension + +``skip`` - Skip an Object + Leaves one object out of a ``prefix`` roll-out. + + .. code-block:: python + + # tvm-ffi-stubgen(skip): my_ffi_extension.Internal + ``import-object`` - Import Object Injects a custom import into generated code. The format is ``;;``. diff --git a/python/tvm_ffi/stub/cli.py b/python/tvm_ffi/stub/cli.py index 3fdffa4b1..dfe217651 100644 --- a/python/tvm_ffi/stub/cli.py +++ b/python/tvm_ffi/stub/cli.py @@ -27,7 +27,7 @@ from typing import TYPE_CHECKING from . import consts as C -from .file_utils import FileInfo, collect_files, syntax_for +from .file_utils import CodeBlock, FileInfo, collect_files, syntax_for from .generator import generator_names, get_generator from .layout import classify, write_coverage_report from .lib_state import ( @@ -82,6 +82,8 @@ def __main__() -> int: print( f'{C.TERM_RED}[Failed] File "{file.path}": {traceback.format_exc()}{C.TERM_RESET}' ) + # Stage 1b. Add the object blocks a `tvm-ffi-stubgen(prefix)` file asks for. + failed += _roll_out_prefixes(files) # Stage 2. Generate stubs if they are not defined on the file. generated_prefixes: set[str] = set() @@ -162,6 +164,64 @@ def _stage_1( ty_map[lhs.strip()] = rhs.strip() +def _roll_out_prefixes(files: list[FileInfo]) -> int: + """Append an ``object/`` block for every registered object under a file's ``prefix``. + + A key with a block in any file of the run, or named by a ``skip`` directive in + the file, is left alone. New blocks follow the file's last object block, parents + first; an ``import-section`` is added after the ``prefix`` line when the file has + none. Returns the number of files whose directives were invalid. + """ + defined = {code.param for file in files for code in file.code_blocks if code.kind == "object"} + registry = collect_type_keys() + owners: dict[str, Path] = {} + failed = 0 + for file in files: + directives = [c for c in file.code_blocks if c.kind == "directive"] + heads = [c for c in directives if c.param[0] == "prefix"] + if not heads: + continue + head = heads[0] + prefix = head.param[1].rstrip(".") + error = "" + if len(heads) > 1: + error = f"more than one `prefix` directive (line {heads[1].lineno_start})" + elif owners.setdefault(prefix, file.path) != file.path: + error = f"prefix `{prefix}` is already declared by {owners[prefix]}" + if error: + failed += 1 + print(f'{C.TERM_RED}[Failed] File "{file.path}": {error}{C.TERM_RESET}') + continue + if prefix not in registry: + print( + f"{C.TERM_YELLOW}[Skipped] No registered object under prefix `{prefix}`{C.TERM_RESET}" + ) + continue + skipped = {c.param[1].strip() for c in directives if c.param[0] == "skip"} + keys = [key for key in registry[prefix] if key not in defined and key not in skipped] + blocks = file.code_blocks + if not any(c.kind == "import-section" for c in blocks): + at = blocks.index(head) + 1 + blocks[at:at] = _new_blocks(file.syntax, head.lineno_start, "import-section") + at = max(i for i, c in enumerate(blocks) if c.kind in ("object", "import-section")) + 1 + blocks[at:at] = [ + block + for info in toposort_objects(keys) + for block in _new_blocks(file.syntax, head.lineno_start, f"object/{info.type_key}") + ] + return failed + + +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}" + block = CodeBlock.from_begin_line(lineno, begin, syntax) + block.lineno_end = lineno + block.lines = [begin, syntax.end] + blank = CodeBlock(kind=None, param="", lineno_start=lineno, lineno_end=lineno, lines=[""]) + return [blank, block] + + def _stage_2( files: list[FileInfo], ty_map: dict[str, str], diff --git a/python/tvm_ffi/stub/consts.py b/python/tvm_ffi/stub/consts.py index ba790681d..5df28d5f8 100644 --- a/python/tvm_ffi/stub/consts.py +++ b/python/tvm_ffi/stub/consts.py @@ -87,7 +87,7 @@ def directive(self, name: str) -> str: #: One-line directive names consumed by the language-neutral pipeline. Generators #: must not declare these names; every other name must be declared by the active #: generator (``Generator.directive_kinds``). -PIPELINE_DIRECTIVE_KINDS: frozenset[str] = frozenset({"ty-map"}) +PIPELINE_DIRECTIVE_KINDS: frozenset[str] = frozenset({"ty-map", "prefix", "skip"}) STUB_BLOCK_KINDS: TypeAlias = Literal[ "global", diff --git a/tests/python/test_stubgen_rust.py b/tests/python/test_stubgen_rust.py index 394934007..dc161cc67 100644 --- a/tests/python/test_stubgen_rust.py +++ b/tests/python/test_stubgen_rust.py @@ -1302,3 +1302,104 @@ def test_stage_3_checks_dependencies_across_the_run(tmp_path: Path) -> None: text = "\n".join(line for block in info.code_blocks for line in block.lines) assert " base: TestCxxClassBaseObj," in text assert "use crate::hand::TestCxxClassBaseObj;" in text + + +# --------------------------------------------------------------------------- +# `prefix` / `skip`: a file declares a namespace and the blocks are rolled out +# --------------------------------------------------------------------------- + + +def test_prefix_directive_rolls_out_a_namespace( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + src = tmp_path / "testing.rs" + src.write_text( + "\n".join( + [ + "//! Hand-written skeleton.", + f"{C.RUST_SYNTAX.directive('prefix')} testing", + f"{C.RUST_SYNTAX.directive('skip')} testing.TestCxxClassDerivedDerived", + "", + f"{C.RUST_SYNTAX.begin} object/testing.TestCxxClassBase", + C.RUST_SYNTAX.end, + "", + "pub fn hand_written() {}", + "", + ] + ), + encoding="utf-8", + ) + check = ["tvm-ffi-stubgen", "--target", "rust", "--check", str(tmp_path)] + monkeypatch.setattr("sys.argv", check) + assert stub_cli.__main__() == 1 # the blocks still to be added make the file stale + monkeypatch.setattr("sys.argv", ["tvm-ffi-stubgen", "--target", "rust", str(tmp_path)]) + assert stub_cli.__main__() == 0 + text = src.read_text(encoding="utf-8") + begins = [line for line in text.splitlines() if line.startswith(C.RUST_SYNTAX.begin)] + # An import section is added after the `prefix` line; the existing block is kept once. + assert begins[0] == f"{C.RUST_SYNTAX.begin} import-section" + assert begins[1] == f"{C.RUST_SYNTAX.begin} object/testing.TestCxxClassBase" + assert text.count("object/testing.TestCxxClassBase\n") == 1 + # The rest of the namespace follows, parents first; the skipped leaf is absent. + derived = begins.index(f"{C.RUST_SYNTAX.begin} object/testing.TestCxxClassDerived") + assert derived > 1 + assert "object/testing.TestCxxClassDerivedDerived" not in text + assert "pub struct TestCxxClassDerivedObj {" in text + assert text.endswith("pub fn hand_written() {}\n") + monkeypatch.setattr("sys.argv", check) + assert stub_cli.__main__() == 0 + + +def test_prefix_declared_twice_is_a_failure( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + for name in ("a.rs", "b.rs"): + (tmp_path / name).write_text( + f"{C.RUST_SYNTAX.directive('prefix')} testing\n", encoding="utf-8" + ) + monkeypatch.setattr( + "sys.argv", ["tvm-ffi-stubgen", "--target", "rust", "--check", str(tmp_path)] + ) + assert stub_cli.__main__() == 2 + + +def test_roll_out_matches_the_prefix_exactly( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr( + stub_cli, "collect_type_keys", lambda: {"a": ["a.Y", "a.X"], "a.b": ["a.b.Z"]} + ) + monkeypatch.setattr( + stub_cli, + "toposort_objects", + lambda keys: [ObjectInfo(fields=[], methods=[], type_key=key) for key in sorted(keys)], + ) + begin, end = C.RUST_SYNTAX.begin, C.RUST_SYNTAX.end + src = tmp_path / "a.rs" + src.write_text( + "\n".join( + [ + f"{C.RUST_SYNTAX.directive('prefix')} a.", + f"{begin} import-section", + end, + "mod tail {}", + "", + ] + ), + encoding="utf-8", + ) + info = FileInfo.from_file(src) + assert info is not None + assert stub_cli._roll_out_prefixes([info]) == 0 + assert [line for block in info.code_blocks for line in block.lines] == [ + f"{C.RUST_SYNTAX.directive('prefix')} a.", + f"{begin} import-section", + end, + "", + f"{begin} object/a.X", + end, + "", + f"{begin} object/a.Y", + end, + "mod tail {}", + ] From b5c4b921429b839704996932cc4c6ae4fdb2b125 Mon Sep 17 00:00:00 2001 From: yuchuan Date: Sat, 5 Sep 2026 21:41:32 -0400 Subject: [PATCH 2/4] simplify. Signed-off-by: yuchuan --- docs/packaging/stubgen.rst | 15 +++------------ python/tvm_ffi/stub/cli.py | 8 +++----- 2 files changed, 6 insertions(+), 17 deletions(-) diff --git a/docs/packaging/stubgen.rst b/docs/packaging/stubgen.rst index a88cf29b8..83b847b17 100644 --- a/docs/packaging/stubgen.rst +++ b/docs/packaging/stubgen.rst @@ -388,22 +388,13 @@ When you run the tool, it: # tvm-ffi-stubgen(ty-map): ffi.reflection.AccessStep -> ffi.access_path.AccessStep ``prefix`` - Demand a Namespace - Declares that the file holds every object registered directly under a type-key - prefix (``tirx`` covers ``tirx.Add`` but not ``tirx.transform.X``). Each run appends - an empty ``object/`` block for every such object that no processed file - defines yet, parents first, after the file's last block (an ``import-section`` is - added after the ``prefix`` line when the file has none), and then fills it. Code - outside the blocks is preserved, so a file can start as a one-line skeleton. + Adds an ``object/`` block to the file for every object registered directly + under the prefix that no processed file defines yet; ``skip`` leaves one out. Code + outside the blocks is preserved. .. code-block:: python # tvm-ffi-stubgen(prefix): my_ffi_extension - -``skip`` - Skip an Object - Leaves one object out of a ``prefix`` roll-out. - - .. code-block:: python - # tvm-ffi-stubgen(skip): my_ffi_extension.Internal ``import-object`` - Import Object diff --git a/python/tvm_ffi/stub/cli.py b/python/tvm_ffi/stub/cli.py index dfe217651..3029d4734 100644 --- a/python/tvm_ffi/stub/cli.py +++ b/python/tvm_ffi/stub/cli.py @@ -165,12 +165,10 @@ def _stage_1( def _roll_out_prefixes(files: list[FileInfo]) -> int: - """Append an ``object/`` block for every registered object under a file's ``prefix``. + """Append an ``object/`` block for each registered object under a file's ``prefix``. - A key with a block in any file of the run, or named by a ``skip`` directive in - the file, is left alone. New blocks follow the file's last object block, parents - first; an ``import-section`` is added after the ``prefix`` line when the file has - none. Returns the number of files whose directives were invalid. + Keys with a block in any file of the run, or named by ``skip``, are left alone; an + ``import-section`` is added when the file has none. 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() From 9ee41f12d87c944e3265af505fc7cec4d5a2876a Mon Sep 17 00:00:00 2001 From: yuchuan Date: Sat, 5 Sep 2026 22:49:54 -0400 Subject: [PATCH 3/4] fix. Signed-off-by: yuchuan --- docs/packaging/stubgen.rst | 12 ++++++------ python/tvm_ffi/stub/cli.py | 6 ++++-- tests/python/test_stubgen_rust.py | 17 +++++++++++++++++ 3 files changed, 27 insertions(+), 8 deletions(-) diff --git a/docs/packaging/stubgen.rst b/docs/packaging/stubgen.rst index 83b847b17..98087e2e9 100644 --- a/docs/packaging/stubgen.rst +++ b/docs/packaging/stubgen.rst @@ -388,14 +388,14 @@ When you run the tool, it: # tvm-ffi-stubgen(ty-map): ffi.reflection.AccessStep -> ffi.access_path.AccessStep ``prefix`` - Demand a Namespace - Adds an ``object/`` block to the file for every object registered directly - under the prefix that no processed file defines yet; ``skip`` leaves one out. Code - outside the blocks is preserved. + Rust target only. Adds an ``object/`` block to the file for every object + registered directly under the prefix that no processed file defines yet; ``skip`` + leaves one out. Code outside the blocks is preserved. - .. code-block:: python + .. code-block:: rust - # tvm-ffi-stubgen(prefix): my_ffi_extension - # tvm-ffi-stubgen(skip): my_ffi_extension.Internal + // tvm-ffi-stubgen(prefix): my_ffi_extension + // tvm-ffi-stubgen(skip): my_ffi_extension.Internal ``import-object`` - Import Object Injects a custom import into generated code. The format is diff --git a/python/tvm_ffi/stub/cli.py b/python/tvm_ffi/stub/cli.py index 3029d4734..ef39f1101 100644 --- a/python/tvm_ffi/stub/cli.py +++ b/python/tvm_ffi/stub/cli.py @@ -82,8 +82,6 @@ def __main__() -> int: print( f'{C.TERM_RED}[Failed] File "{file.path}": {traceback.format_exc()}{C.TERM_RESET}' ) - # Stage 1b. Add the object blocks a `tvm-ffi-stubgen(prefix)` file asks for. - failed += _roll_out_prefixes(files) # Stage 2. Generate stubs if they are not defined on the file. generated_prefixes: set[str] = set() @@ -98,6 +96,10 @@ def __main__() -> int: generator=generator, ) + # Stage 2b. Add the object blocks a `tvm-ffi-stubgen(prefix)` file asks for. This runs + # after `--init`, which rewrites files on disk and reloads them. + failed += _roll_out_prefixes(files) + # Stage 3: Process # - `tvm-ffi-stubgen(begin): global/...` # - `tvm-ffi-stubgen(begin): object/...` diff --git a/tests/python/test_stubgen_rust.py b/tests/python/test_stubgen_rust.py index dc161cc67..994fa9e88 100644 --- a/tests/python/test_stubgen_rust.py +++ b/tests/python/test_stubgen_rust.py @@ -1403,3 +1403,20 @@ def test_roll_out_matches_the_prefix_exactly( end, "mod tail {}", ] + + +def test_prefix_survives_init(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + """`--init` rewrites the file on disk and reloads it; the roll-out must come after that.""" + (tmp_path / "testing").mkdir() + mod_rs = tmp_path / "testing" / "mod.rs" + mod_rs.write_text(f"{C.RUST_SYNTAX.directive('prefix')} testing\n", encoding="utf-8") + init = ["--init-pypkg", "demo", "--init-lib", "demo_shared", "--init-prefix", "testing."] + monkeypatch.setattr("sys.argv", ["tvm-ffi-stubgen", "--target", "rust", *init, str(tmp_path)]) + assert stub_cli.__main__() == 0 + text = mod_rs.read_text(encoding="utf-8") + assert text.startswith(f"{C.RUST_SYNTAX.directive('prefix')} testing\n") + assert "pub struct TestCxxClassDerivedObj {" in text + monkeypatch.setattr( + "sys.argv", ["tvm-ffi-stubgen", "--target", "rust", "--check", str(tmp_path)] + ) + assert stub_cli.__main__() == 0 From 897f961dffeb112a09604dccd015d2084d02df54 Mon Sep 17 00:00:00 2001 From: yuchuan Date: Sat, 5 Sep 2026 23:03:37 -0400 Subject: [PATCH 4/4] fix. Signed-off-by: yuchuan --- python/tvm_ffi/stub/cli.py | 12 ++++++++++-- tests/python/test_stubgen_rust.py | 14 ++++++++++++-- 2 files changed, 22 insertions(+), 4 deletions(-) diff --git a/python/tvm_ffi/stub/cli.py b/python/tvm_ffi/stub/cli.py index ef39f1101..4f96e9aed 100644 --- a/python/tvm_ffi/stub/cli.py +++ b/python/tvm_ffi/stub/cli.py @@ -172,7 +172,9 @@ def _roll_out_prefixes(files: list[FileInfo]) -> int: Keys with a block in any file of the run, or named by ``skip``, are left alone; an ``import-section`` is added when the file has none. Returns the number of bad files. """ - defined = {code.param for file in files for code in file.code_blocks if code.kind == "object"} + defined = { + code.param for file in files for code in file.code_blocks if code.kind == "object" + } | C.BUILTIN_TYPE_KEYS registry = collect_type_keys() owners: dict[str, Path] = {} failed = 0 @@ -250,6 +252,12 @@ def _find_or_insert_file(path: Path) -> FileInfo: defined_objs: set[str] = { # ty: ignore[invalid-assignment] code.param for file in files for code in file.code_blocks if code.kind == "object" } | C.BUILTIN_TYPE_KEYS + skipped: set[str] = { + code.param[1].strip() + for file in files + for code in file.code_blocks + if code.kind == "directive" and code.param[0] == "skip" + } # Step 0. Generate missing `_ffi_api.py` and `__init__.py` under each prefix. prefix_filter = init_cfg.prefix.strip() @@ -267,7 +275,7 @@ def _find_or_insert_file(path: Path) -> FileInfo: [] if prefix in defined_func_prefixes else global_funcs.get(prefix, []), key=lambda f: f.schema.name, ) - objs = sorted(set(obj_names) - defined_objs) + objs = sorted(set(obj_names) - defined_objs - skipped) object_infos = toposort_objects(objs) if not funcs and not object_infos: continue diff --git a/tests/python/test_stubgen_rust.py b/tests/python/test_stubgen_rust.py index 994fa9e88..97a446d21 100644 --- a/tests/python/test_stubgen_rust.py +++ b/tests/python/test_stubgen_rust.py @@ -1367,7 +1367,9 @@ def test_roll_out_matches_the_prefix_exactly( tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> None: monkeypatch.setattr( - stub_cli, "collect_type_keys", lambda: {"a": ["a.Y", "a.X"], "a.b": ["a.b.Z"]} + stub_cli, + "collect_type_keys", + lambda: {"a": ["a.Y", "a.X"], "a.b": ["a.b.Z"], "ffi": ["ffi.Object"]}, ) monkeypatch.setattr( stub_cli, @@ -1403,19 +1405,27 @@ def test_roll_out_matches_the_prefix_exactly( end, "mod tail {}", ] + # Builtin type keys live in the crate and are never rolled out. + src.write_text(f"{C.RUST_SYNTAX.directive('prefix')} ffi\n", encoding="utf-8") + info = FileInfo.from_file(src) + assert info is not None + assert stub_cli._roll_out_prefixes([info]) == 0 + assert not any(block.kind == "object" for block in info.code_blocks) def test_prefix_survives_init(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: """`--init` rewrites the file on disk and reloads it; the roll-out must come after that.""" (tmp_path / "testing").mkdir() mod_rs = tmp_path / "testing" / "mod.rs" - mod_rs.write_text(f"{C.RUST_SYNTAX.directive('prefix')} testing\n", encoding="utf-8") + skip = f"{C.RUST_SYNTAX.directive('skip')} testing.TestCxxClassDerivedDerived" + mod_rs.write_text(f"{C.RUST_SYNTAX.directive('prefix')} testing\n{skip}\n", encoding="utf-8") init = ["--init-pypkg", "demo", "--init-lib", "demo_shared", "--init-prefix", "testing."] monkeypatch.setattr("sys.argv", ["tvm-ffi-stubgen", "--target", "rust", *init, str(tmp_path)]) assert stub_cli.__main__() == 0 text = mod_rs.read_text(encoding="utf-8") assert text.startswith(f"{C.RUST_SYNTAX.directive('prefix')} testing\n") assert "pub struct TestCxxClassDerivedObj {" in text + assert "object/testing.TestCxxClassDerivedDerived" not in text # `--init` honours `skip` monkeypatch.setattr( "sys.argv", ["tvm-ffi-stubgen", "--target", "rust", "--check", str(tmp_path)] )