[STUBGEN] Let a file declare a namespace with a prefix directive

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>``.

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/<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 _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

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",

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