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 4e754f9f [STUBGEN] Let a file declare a namespace with a `prefix`
directive (#752)
4e754f9f is described below
commit 4e754f9f647e16ffe0aee860531f55d62e33138b
Author: Linzhang Li <[email protected]>
AuthorDate: Sun Sep 6 08:06:42 2026 -0400
[STUBGEN] Let a file declare a namespace with a `prefix` directive (#752)
## Summary
A file that starts with
```rust
// tvm-ffi-stubgen(prefix): tirx
// tvm-ffi-stubgen(skip): tirx.DispatchContext
```
owns every object registered directly under `tirx` (`tirx.transform.*`
is another file). On each run the tool appends an empty `object/<key>`
block for every such object that no file of the run defines yet, parents
first, after the file's last block, and fills it as usual; an
`import-section` is added after the `prefix` line when the file has
none. `skip` leaves one object out, and a skipped object counts as not
asked for, so the dependency error from #748 names it if something else
needs it. Code outside the blocks is untouched, so a one-line skeleton
is a valid starting point. Rust target only: a Python `object/` block
needs a class around it.
The roll-out happens in memory before stage 3, so the new blocks join
`declared` and `--check` (#746) reports a file with pending blocks as
stale. It runs after `--init`, which rewrites files on disk; `--init`
honours `skip` as well. Builtin `ffi.*` keys are never rolled out. Two
`prefix` lines in one file, or one prefix declared by two files, fail
the run (exit 2); an unknown prefix is only a warning.
## Changes
- `stub/cli.py`: `_roll_out_prefixes` and `_new_blocks`, called once
after `--init`, returning a failure count; `_stage_2` subtracts the
run's `skip` keys.
- `stub/consts.py`: `prefix` and `skip` join the pipeline directive
kinds.
- `docs/packaging/stubgen.rst`: the directive in the reference.
- `tests/python/test_stubgen_rust.py`: the `testing` namespace rolled
out around an existing block, a skip, and a hand-written tail (`--check`
1, in place 0, `--check` 0); the same skeleton under `--init` (blocks
written, skip honoured, `--check` 0); a prefix declared twice; exact
matching, insertion order, and the builtin exclusion with a stubbed
registry.
## Testing
`test_stubgen_rust.py` + `test_stubgen.py`: 109 passed; pre-commit lint
clean. Against `libtvm_compiler.so`, a `tirx.rs` with the two lines
above, one more `skip`, and thirteen `ty-map` lines for the `ir` types
rolls out 63 blocks in one run, leaves no `super::` reference, keeps the
trailing `mod tail {}`, and passes `--check` on the next run.
---------
Signed-off-by: yuchuan <[email protected]>
---
docs/packaging/stubgen.rst | 10 +++
python/tvm_ffi/stub/cli.py | 72 ++++++++++++++++++++-
python/tvm_ffi/stub/consts.py | 2 +-
tests/python/test_stubgen_rust.py | 128 ++++++++++++++++++++++++++++++++++++++
4 files changed, 209 insertions(+), 3 deletions(-)
diff --git a/docs/packaging/stubgen.rst b/docs/packaging/stubgen.rst
index 3b446a9d..98087e2e 100644
--- a/docs/packaging/stubgen.rst
+++ b/docs/packaging/stubgen.rst
@@ -387,6 +387,16 @@ When you run the tool, it:
# tvm-ffi-stubgen(ty-map): ffi.reflection.AccessStep ->
ffi.access_path.AccessStep
+``prefix`` - Demand a Namespace
+ 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. Code outside the blocks is preserved.
+
+ .. code-block:: rust
+
+ // 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
``<full_name>;<type_checking_only>;<alias>``.
diff --git a/python/tvm_ffi/stub/cli.py b/python/tvm_ffi/stub/cli.py
index 3fdffa4b..4f96e9ae 100644
--- a/python/tvm_ffi/stub/cli.py
+++ b/python/tvm_ffi/stub/cli.py
@@ -27,7 +27,7 @@ from pathlib import Path
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 (
@@ -96,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/...`
@@ -162,6 +166,64 @@ def _stage_1(
ty_map[lhs.strip()] = rhs.strip()
+def _roll_out_prefixes(files: list[FileInfo]) -> 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, 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"
+ } | C.BUILTIN_TYPE_KEYS
+ 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],
@@ -190,6 +252,12 @@ def _stage_2(
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()
@@ -207,7 +275,7 @@ def _stage_2(
[] 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/python/tvm_ffi/stub/consts.py b/python/tvm_ffi/stub/consts.py
index ba790681..5df28d5f 100644
--- a/python/tvm_ffi/stub/consts.py
+++ b/python/tvm_ffi/stub/consts.py
@@ -87,7 +87,7 @@ SYNTAX_BY_EXT: dict[str, MarkerSyntax] = {
#: 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 39493400..97a446d2 100644
--- a/tests/python/test_stubgen_rust.py
+++ b/tests/python/test_stubgen_rust.py
@@ -1302,3 +1302,131 @@ 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"], "ffi": ["ffi.Object"]},
+ )
+ 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 {}",
+ ]
+ # 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"
+ 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)]
+ )
+ assert stub_cli.__main__() == 0
