diff --git a/docs/packaging/stubgen.rst b/docs/packaging/stubgen.rst index 3b446a9d0..98087e2e9 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/`` 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 ``;;``. diff --git a/python/tvm_ffi/stub/cli.py b/python/tvm_ffi/stub/cli.py index 3fdffa4b1..4f96e9aed 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 ( @@ -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/`` 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 _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() @@ -207,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/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..97a446d21 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