This is an automated email from the ASF dual-hosted git repository.
tlopex 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 ef5ce4e0 [STUBGEN] Generalize one-line directive handling (#736)
ef5ce4e0 is described below
commit ef5ce4e0c1700e028545d49a010b1689ed808234
Author: Linzhang Li <[email protected]>
AuthorDate: Thu Sep 3 14:39:02 2026 -0400
[STUBGEN] Generalize one-line directive handling (#736)
## Summary
Generalize parsing and dispatch for standalone one-line
`tvm-ffi-stubgen`
directives. Pipeline-owned and generator-owned directives now share one
`CodeBlock` representation, while each layer remains responsible for its
own
payload grammar. This lets a target generator add directives without
extending
the language-neutral file parser.
## Changes
- Add `MarkerSyntax.directive(name)` for constructing standalone
directive
markers with either Python or Rust comment syntax.
- Parse standalone directives uniformly as
`CodeBlock(kind="directive", param=(name, payload))` and preserve their
source
lines when rewriting files.
- Add `PIPELINE_DIRECTIVE_KINDS` for directives consumed by the common
pipeline.
`ty-map` remains pipeline-owned and is processed before init-mode
generation.
- Add `Generator.directive_kinds` and `Generator.add_directive` so each
target
declares and handles its own directive names. Undeclared names fail with
their
source line number.
- Route the existing Python `import-object` support through
`add_directive`, and
move its semicolon-delimited payload parsing from the common parser into
the
Python generator.
## Compatibility
The undocumented `tvm-ffi-stubgen(begin): ty-map/...` block form is no
longer
accepted; `ty-map` is a standalone one-line directive. Out-of-tree
`Generator`
implementations must provide `directive_kinds` and `add_directive`;
`add_imported_object` is no longer called by the pipeline.
## Tests
- `uv run pytest -q tests/python/test_stubgen.py` (43 passed)
- `uv run ruff check` on all changed Python files
- `git diff --check`
---------
Signed-off-by: yuchuan <[email protected]>
---
python/tvm_ffi/stub/cli.py | 31 +++---
python/tvm_ffi/stub/consts.py | 20 ++--
python/tvm_ffi/stub/file_utils.py | 51 +++-------
python/tvm_ffi/stub/generator.py | 23 +++--
python/tvm_ffi/stub/python_generator/generator.py | 14 +--
tests/python/test_stubgen.py | 118 ++++++++++++++++++++--
6 files changed, 180 insertions(+), 77 deletions(-)
diff --git a/python/tvm_ffi/stub/cli.py b/python/tvm_ffi/stub/cli.py
index 2b054c72..3dfc8151 100644
--- a/python/tvm_ffi/stub/cli.py
+++ b/python/tvm_ffi/stub/cli.py
@@ -120,15 +120,15 @@ def _stage_1(
ty_map: dict[str, str],
) -> None:
for code in file.code_blocks:
- if code.kind == "ty-map":
- try:
- assert isinstance(code.param, str)
- lhs, rhs = code.param.split("->")
- except ValueError as e:
- raise ValueError(
- f"Invalid ty_map format at line {code.lineno_start}.
Example: `A.B -> C.D`"
- ) from e
- ty_map[lhs.strip()] = rhs.strip()
+ if code.kind != "directive" or code.param[0] != "ty-map":
+ continue
+ try:
+ lhs, rhs = code.param[1].split("->")
+ except ValueError as e:
+ raise ValueError(
+ f"Invalid ty_map format at line {code.lineno_start}. Example:
`A.B -> C.D`"
+ ) from e
+ ty_map[lhs.strip()] = rhs.strip()
def _stage_2(
@@ -221,11 +221,16 @@ def _stage_3( # noqa: PLR0912
defined_funcs: set[str] = set()
defined_types: set[str] = set()
imports = generator.new_imports()
- # Stage 1. Collect `tvm-ffi-stubgen(import-object): ...`
+ # Stage 1. Hand the one-line directives the pipeline does not consume
itself to the generator.
for code in file.code_blocks:
- if code.kind == "import-object":
- name, type_checking_only, alias = code.param
- generator.add_imported_object(imports, name, type_checking_only,
alias)
+ if code.kind != "directive":
+ continue
+ name, payload = code.param
+ if name in C.PIPELINE_DIRECTIVE_KINDS:
+ continue # consumed by `_stage_1`
+ if name not in generator.directive_kinds:
+ raise ValueError(f"Unknown directive `{name}` at line
{code.lineno_start}")
+ generator.add_directive(imports, name, payload, code.lineno_start)
# Stage 2. Process `tvm-ffi-stubgen(begin): global/...`
for code in file.code_blocks:
if code.kind == "global":
diff --git a/python/tvm_ffi/stub/consts.py b/python/tvm_ffi/stub/consts.py
index 53115386..ba790681 100644
--- a/python/tvm_ffi/stub/consts.py
+++ b/python/tvm_ffi/stub/consts.py
@@ -54,19 +54,23 @@ class MarkerSyntax:
@property
def ty_map(self) -> str:
- """One-line type-map directive: ``<comment>
tvm-ffi-stubgen(ty-map):``."""
- return f"{self.prefix}ty-map):"
+ """The ``ty-map`` directive marker: ``<comment>
tvm-ffi-stubgen(ty-map):``."""
+ return self.directive("ty-map")
@property
def import_object(self) -> str:
- """One-line import-object directive: ``<comment>
tvm-ffi-stubgen(import-object):``."""
- return f"{self.prefix}import-object):"
+ """The ``import-object`` directive marker: ``<comment>
tvm-ffi-stubgen(import-object):``."""
+ return self.directive("import-object")
@property
def skip_file(self) -> str:
"""Whole-file opt-out directive: ``<comment>
tvm-ffi-stubgen(skip-file)``."""
return f"{self.prefix}skip-file)"
+ def directive(self, name: str) -> str:
+ """One-line directive marker: ``<comment>
tvm-ffi-stubgen(<name>):``."""
+ return f"{self.prefix}{name}):"
+
PYTHON_SYNTAX = MarkerSyntax(comment="#")
RUST_SYNTAX = MarkerSyntax(comment="//")
@@ -80,12 +84,16 @@ SYNTAX_BY_EXT: dict[str, MarkerSyntax] = {
".rs": RUST_SYNTAX,
}
+#: 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"})
+
STUB_BLOCK_KINDS: TypeAlias = Literal[
"global",
"object",
- "ty-map",
"import-section",
- "import-object",
+ "directive",
"export",
"__all__",
None,
diff --git a/python/tvm_ffi/stub/file_utils.py
b/python/tvm_ffi/stub/file_utils.py
index 94d0769c..381254cf 100644
--- a/python/tvm_ffi/stub/file_utils.py
+++ b/python/tvm_ffi/stub/file_utils.py
@@ -49,9 +49,8 @@ class CodeBlock:
assert self.kind in {
"global",
"object",
- "ty-map",
"import-section",
- "import-object",
+ "directive",
"export",
"__all__",
None,
@@ -68,23 +67,15 @@ class CodeBlock:
@staticmethod
def from_begin_line(lineo: int, line: str, syntax: C.MarkerSyntax) ->
CodeBlock:
"""Parse a line to create a CodeBlock if it contains a stub begin
marker."""
- if line.startswith(syntax.ty_map):
- line = line[len(syntax.ty_map) :].strip()
+ if not line.startswith(syntax.begin):
+ # One-line directive `<comment> tvm-ffi-stubgen(<name>):
<payload>`. The payload is
+ # kept verbatim for whoever consumes the name: the pipeline or the
generator.
+ name, sep, payload = line[len(syntax.prefix) :].partition("):")
+ if not sep or not name or not name.replace("-",
"_").isidentifier():
+ raise ValueError(f"Unknown stub type at line {lineo}: {line}")
return CodeBlock(
- kind="ty-map",
- param=line,
- lineno_start=lineo,
- lineno_end=lineo,
- lines=[],
- )
- elif line.startswith(syntax.import_object):
- line = line[len(syntax.import_object) :].strip()
- splits = [p.strip() for p in line.split(";")]
- if len(splits) < 3:
- splits += [""] * (3 - len(splits))
- return CodeBlock(
- kind="import-object",
- param=tuple(splits),
+ kind="directive",
+ param=(name, payload.strip()),
lineno_start=lineo,
lineno_end=lineo,
lines=[],
@@ -99,9 +90,6 @@ class CodeBlock:
elif stub.startswith("object/"):
kind = "object"
param = stub[len("object/") :].strip()
- elif stub.startswith("ty-map/"):
- kind = "ty-map"
- param = stub[len("ty-map/") :].strip()
elif stub == "import-section":
kind = "import-section"
param = ""
@@ -199,22 +187,13 @@ class FileInfo:
code.lines.append(line)
codes.append(code)
code = None
- elif clean_line.startswith(syntax.ty_map):
- # Process "<comment> tvm-ffi-stubgen(ty_map)"
- ty_code = CodeBlock.from_begin_line(lineno, clean_line, syntax)
- ty_code.lineno_end = lineno
- ty_code.lines.append(line)
- codes.append(ty_code)
- del ty_code
- elif clean_line.startswith(syntax.import_object):
- # Process "<comment> tvm-ffi-stubgen(import-object)"
- imp_code = CodeBlock.from_begin_line(lineno, clean_line,
syntax)
- imp_code.lineno_end = lineno
- imp_code.lines.append(line)
- codes.append(imp_code)
- del imp_code
elif clean_line.startswith(syntax.prefix):
- raise ValueError(f"Unknown stub type at line {lineno}:
{clean_line}")
+ # Process a one-line directive "<comment>
tvm-ffi-stubgen(<name>): ..."
+ dir_code = CodeBlock.from_begin_line(lineno, clean_line,
syntax)
+ dir_code.lineno_end = lineno
+ dir_code.lines.append(line)
+ codes.append(dir_code)
+ del dir_code
elif code is None:
# Process a plain line outside of any stub block
codes.append(
diff --git a/python/tvm_ffi/stub/generator.py b/python/tvm_ffi/stub/generator.py
index 99c3aa0d..3cd896a3 100644
--- a/python/tvm_ffi/stub/generator.py
+++ b/python/tvm_ffi/stub/generator.py
@@ -30,9 +30,10 @@ The stub generator separates two concerns:
A :class:`Generator` encapsulates concern (2); ``cli.py`` drives concern (1)
and
delegates every act of emitting text — and every act of collecting imports — to
the active generator. The import collector is opaque to the pipeline:
``cli.py``
-asks the generator to create one, seed it from ``import-object`` directives,
and
-later render it, but never reaches inside. Adding a language is therefore
-"implement one more :class:`Generator`" rather than forking the pipeline.
+asks the generator to create one, seed it from the one-line directives the
+generator declares, and later render it, but never reaches inside. Adding a
+language is therefore "implement one more :class:`Generator`" rather than
+forking the pipeline.
"""
from __future__ import annotations
@@ -74,6 +75,11 @@ class Generator(Protocol):
#: Comment-marker syntax for the files this generator emits.
syntax: C.MarkerSyntax
+ #: Names of the one-line directives (``<comment> tvm-ffi-stubgen(<name>):
<payload>``)
+ #: this generator consumes. Names in
:data:`consts.PIPELINE_DIRECTIVE_KINDS` belong
+ #: to the pipeline; any other undeclared name is an error.
+ directive_kinds: frozenset[str]
+
def default_ty_map(self) -> dict[str, str]:
"""Return the default FFI-origin -> target-type name map for this
language."""
...
@@ -84,10 +90,13 @@ class Generator(Protocol):
"""Create a fresh, empty import collector for one file."""
...
- def add_imported_object(
- self, imports: Any, name: str, type_checking_only: str, alias: str
- ) -> None:
- """Record an ``import-object`` directive (raw directive fields) into
``imports``."""
+ def add_directive(self, imports: Any, name: str, payload: str, lineno:
int) -> None:
+ """Record a one-line directive (raw payload) into ``imports``.
+
+ The collector is per file, so a directive applies to the blocks of the
+ file it appears in. ``name`` is always one of :attr:`directive_kinds`;
+ the payload's grammar is the generator's to define.
+ """
...
def canonical_type_name(self, type_key: str) -> str:
diff --git a/python/tvm_ffi/stub/python_generator/generator.py
b/python/tvm_ffi/stub/python_generator/generator.py
index d05b2f2a..906c9a8f 100644
--- a/python/tvm_ffi/stub/python_generator/generator.py
+++ b/python/tvm_ffi/stub/python_generator/generator.py
@@ -44,6 +44,7 @@ class PythonGenerator:
name = "python"
syntax = C.PYTHON_SYNTAX
source_exts = frozenset({".py", ".pyi"})
+ directive_kinds: frozenset[str] = frozenset({"import-object"})
def default_ty_map(self) -> dict[str, str]:
"""Return the default FFI-origin -> Python-type name map."""
@@ -55,13 +56,14 @@ class PythonGenerator:
"""Create an empty import collector."""
return PythonImports()
- def add_imported_object(
- self, imports: PythonImports, name: str, type_checking_only: str,
alias: str
- ) -> None:
- """Record an ``import-object`` directive into the collector."""
+ def add_directive(self, imports: PythonImports, name: str, payload: str,
lineno: int) -> None:
+ """Record an ``import-object`` directive
(``<full_name>;<type_checking_only>;<alias>``)."""
+ assert name == "import-object", name
+ parts = [part.strip() for part in payload.split(";")]
+ full_name, type_checking_only, alias = parts + [""] * (3 - len(parts))
tco = type_checking_only.lower() == "true"
- imports.items.append(ImportItem(name, type_checking_only=tco,
alias=alias or None))
- if alias == "_FFI_LOAD_LIB" or
name.endswith("libinfo.load_lib_module"):
+ imports.items.append(ImportItem(full_name, type_checking_only=tco,
alias=alias or None))
+ if alias == "_FFI_LOAD_LIB" or
full_name.endswith("libinfo.load_lib_module"):
imports.has_lib_load = True
def canonical_type_name(self, type_key: str) -> str:
diff --git a/tests/python/test_stubgen.py b/tests/python/test_stubgen.py
index 199ad796..c2606e46 100644
--- a/tests/python/test_stubgen.py
+++ b/tests/python/test_stubgen.py
@@ -26,9 +26,10 @@ from tvm_ffi import Object, method
from tvm_ffi.core import MISSING, TypeSchema,
_lookup_or_register_type_info_from_type_key
from tvm_ffi.dataclasses import py_class
from tvm_ffi.stub import consts as C
-from tvm_ffi.stub.cli import _stage_2, _stage_3
+from tvm_ffi.stub.cli import _stage_1, _stage_2, _stage_3
from tvm_ffi.stub.file_utils import CodeBlock, FileInfo, collect_files,
syntax_for
from tvm_ffi.stub.generator import generator_names, get_generator
+from tvm_ffi.stub.python_generator import PythonGenerator
from tvm_ffi.stub.python_generator import consts as PC
from tvm_ffi.stub.python_generator.codegen import (
generate_python_all,
@@ -88,7 +89,6 @@ def test_codeblock_from_begin_line_variants() -> None:
(f"{C.PYTHON_SYNTAX.begin} global/demo", "global", ("demo", "")),
(f"{C.PYTHON_SYNTAX.begin} global/[email protected]", "global", ("demo",
".registry")),
(f"{C.PYTHON_SYNTAX.begin} object/demo.TypeBase", "object",
"demo.TypeBase"),
- (f"{C.PYTHON_SYNTAX.begin} ty-map/custom", "ty-map", "custom"),
(f"{C.PYTHON_SYNTAX.begin} import-section", "import-section", ""),
]
for lineno, (line, kind, param) in enumerate(cases, start=1):
@@ -103,8 +103,8 @@ def test_codeblock_from_begin_line_variants() -> None:
def test_codeblock_from_begin_line_ty_map_and_unknown() -> None:
line = f"{C.PYTHON_SYNTAX.ty_map} custom -> mapped"
block = CodeBlock.from_begin_line(5, line, C.PYTHON_SYNTAX)
- assert block.kind == "ty-map"
- assert block.param == "custom -> mapped"
+ assert block.kind == "directive"
+ assert block.param == ("ty-map", "custom -> mapped")
assert block.lineno_start == 5
assert block.lineno_end == 5
@@ -153,8 +153,8 @@ def test_fileinfo_from_file_parses_blocks(tmp_path: Path)
-> None:
C.PYTHON_SYNTAX.end,
]
- assert ty_map.kind == "ty-map"
- assert ty_map.param == "x -> y"
+ assert ty_map.kind == "directive"
+ assert ty_map.param == ("ty-map", "x -> y")
assert ty_map.lineno_start == ty_map.lineno_end == 5
assert ty_map.lines == [f"{C.PYTHON_SYNTAX.ty_map} x -> y"]
@@ -804,8 +804,8 @@ def test_stage_3_adds_LIB_when_load_lib_imported(tmp_path:
Path) -> None:
lines=[f"{C.PYTHON_SYNTAX.begin} global/testing", C.PYTHON_SYNTAX.end],
)
import_obj_block = CodeBlock(
- kind="import-object",
- param=("tvm_ffi.libinfo.load_lib_module", "False", "_FFI_LOAD_LIB"),
+ kind="directive",
+ param=("import-object",
"tvm_ffi.libinfo.load_lib_module;False;_FFI_LOAD_LIB"),
lineno_start=1,
lineno_end=1,
lines=[
@@ -1075,8 +1075,9 @@ def test_rust_marker_syntax_parses_rs_file(tmp_path:
Path) -> None:
info = FileInfo.from_file(rs)
assert info is not None
assert info.syntax is C.RUST_SYNTAX
- assert [block.kind for block in info.code_blocks] == ["object", "ty-map"]
+ assert [block.kind for block in info.code_blocks] == ["object",
"directive"]
assert info.code_blocks[0].param == "demo.Foo"
+ assert info.code_blocks[1].param == ("ty-map", "a -> b")
assert info.code_blocks[0].lines[1] == "pub struct Foo;"
# Python markers are plain text inside a Rust file.
@@ -1107,3 +1108,102 @@ def test_generator_registry_names() -> None:
assert generator_names() == ["python"]
with pytest.raises(ValueError, match="Known generators: python"):
get_generator("rust")
+
+
+def test_codeblock_from_begin_line_directive() -> None:
+ """``<comment> tvm-ffi-stubgen(<name>): <payload>`` parses into a one-line
directive block."""
+ for syntax in (C.PYTHON_SYNTAX, C.RUST_SYNTAX):
+ block = CodeBlock.from_begin_line(7, f"{syntax.directive('field')}
a.B.c -> D ", syntax)
+ assert block.kind == "directive"
+ assert block.param == ("field", "a.B.c -> D")
+ assert (block.lineno_start, block.lineno_end) == (7, 7)
+ assert block.lines == []
+ empty = CodeBlock.from_begin_line(1, C.PYTHON_SYNTAX.directive("opaque"),
C.PYTHON_SYNTAX)
+ assert empty.param == ("opaque", "")
+
+ # The `name):` shape is required; anything else stays an unknown marker.
+ for line in (f"{C.PYTHON_SYNTAX.prefix}field)",
f"{C.PYTHON_SYNTAX.prefix}two words): x"):
+ with pytest.raises(ValueError, match="Unknown stub type at line 3"):
+ CodeBlock.from_begin_line(3, line, C.PYTHON_SYNTAX)
+
+
+def test_fileinfo_from_file_keeps_directive_lines(tmp_path: Path) -> None:
+ """A directive is parsed as its own block and written back untouched."""
+ src = tmp_path / "mod.rs"
+ lines = [
+ " // tvm-ffi-stubgen(nullable): ir.Expr.span",
+ f"{C.RUST_SYNTAX.begin} object/demo.Foo",
+ C.RUST_SYNTAX.end,
+ f"{C.RUST_SYNTAX.prefix}bogus)",
+ ]
+ src.write_text("\n".join(lines) + "\n", encoding="utf-8")
+ with pytest.raises(ValueError, match="Unknown stub type at line 4"):
+ FileInfo.from_file(src)
+
+ src.write_text("\n".join(lines[:3]) + "\n", encoding="utf-8")
+ info = FileInfo.from_file(src)
+ assert info is not None
+ assert [block.kind for block in info.code_blocks] == ["directive",
"object"]
+ assert info.code_blocks[0].param == ("nullable", "ir.Expr.span")
+ assert info.code_blocks[0].lines == [lines[0]]
+ assert info.update(verbose=False, dry_run=False) is False
+ assert src.read_text(encoding="utf-8") == "\n".join(lines[:3]) + "\n"
+
+
+def test_stage_3_routes_directives_by_generator_declaration(tmp_path: Path) ->
None:
+ """Pipeline-owned names stay in the pipeline; declared names reach the
generator."""
+ path = tmp_path / "demo.py"
+ ty_map_directive = CodeBlock(
+ kind="directive",
+ param=("ty-map", "demo.Foo -> mapped.Foo"),
+ lineno_start=1,
+ lineno_end=1,
+ lines=[f"{C.PYTHON_SYNTAX.ty_map} demo.Foo -> mapped.Foo"],
+ )
+ directive = CodeBlock(
+ kind="directive",
+ param=("field", "demo.Foo.x -> Bar"),
+ lineno_start=2,
+ lineno_end=2,
+ lines=[f"{C.PYTHON_SYNTAX.directive('field')} demo.Foo.x -> Bar"],
+ )
+ all_block = CodeBlock(
+ kind="__all__",
+ param="",
+ lineno_start=3,
+ lineno_end=4,
+ lines=[f"{C.PYTHON_SYNTAX.begin} __all__", C.PYTHON_SYNTAX.end],
+ )
+ blocks = (ty_map_directive, directive, all_block)
+
+ def _file_info() -> FileInfo:
+ return FileInfo(
+ path=path,
+ lines=tuple(line for block in blocks for line in block.lines),
+ code_blocks=list(blocks),
+ syntax=C.PYTHON_SYNTAX,
+ )
+
+ # `ty-map` is pipeline-owned: `_stage_1` consumes it and `_stage_3` skips
it.
+ assert "ty-map" in C.PIPELINE_DIRECTIVE_KINDS
+ ty_map: dict[str, str] = {}
+ _stage_1(_file_info(), ty_map)
+ assert ty_map == {"demo.Foo": "mapped.Foo"}
+
+ # The Python generator does not declare `field`, so the name is unknown to
it.
+ python = get_generator("python")
+ assert python.directive_kinds == frozenset({"import-object"})
+ with pytest.raises(ValueError, match="Unknown directive `field` at line
2"):
+ _stage_3(_file_info(), Options(dry_run=True), _default_ty_map(), {},
python)
+
+ seen: list[tuple[object, str, str, int]] = []
+
+ class FieldAware(PythonGenerator):
+ directive_kinds = frozenset({"field"})
+
+ def add_directive(self, imports: object, name: str, payload: str,
lineno: int) -> None:
+ seen.append((imports, name, payload, lineno))
+
+ _stage_3(_file_info(), Options(dry_run=True), _default_ty_map(), {},
FieldAware())
+ assert [entry[1:] for entry in seen] == [("field", "demo.Foo.x -> Bar", 2)]
+ assert isinstance(seen[0][0], type(python.new_imports()))
