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 e2179b2e [Stubgen] One grammar for one-line directives: `subject:
target(keys)` (#968)
e2179b2e is described below
commit e2179b2e82d39be91076da07d840720b6e706359
Author: Tianqi Chen <[email protected]>
AuthorDate: Sun Nov 8 16:50:53 2026 -0500
[Stubgen] One grammar for one-line directives: `subject: target(keys)`
(#968)
Stubgen's one-line directives each had their own payload syntax (`a ->
b`, `T.f -> N(i32) { A=0 }`, `x;False;y`, bare `allow`). This replaces
all of them with one grammar, shared by the Python and Rust generators:
```
<comment> tvm-ffi-stubgen(<subject>): <target>(<key>, <key> = <value>, ...)
```
The parentheses are omitted when there are no keys. Lines for the same
subject and target merge; a repeatable key (`upcast`) collects every
value, and any other key may repeat only with the same value. Commas and
`=` inside brackets belong to the value, so `pub(crate)` and `Option<A,
B>` stay whole.
| Before | After |
|---|---|
| `ty-map: A -> B` | `object: A(ty_map = B)` |
| `skip: T` / `opaque: T` / `custom-new: T` / `no-new: T` | `object:
T(skip)` / `(opaque)` / `(custom_new)` / `(no_new)` |
| `upcast: T -> V` | `object: T(upcast = V)` |
| `field: T.f -> Ty` / `nullable: T.f` / `visibility: T.f -> pub(crate)`
| `field: T.f(type = Ty, nullable, visibility = pub(crate))` |
| `enum: T.f -> K(i32) closed { A=0, B=1 }` | `enum: K(repr = i32,
closed, A = 0, B = 1)` + `field: T.f(type = K)` |
| `prefix: p` / `missing-docs: allow` | `file: mod(prefix = p,
missing_docs = allow)` |
| `import-object: a.b.C;True;_C` | `import: a.b.C(type_checking_only,
alias = _C)` |
`begin`/`end` blocks and `skip-file` are unchanged.
A small parser (`stub/directives.py`) checks each key against a
per-subject table. The pipeline declares `object(skip, ty_map)` and
`file(prefix)`. Each generator declares its own keys in
`Generator.directive_keys` and gets the parsed file through
`new_imports`. In-tree directives, the `rust_stubgen` example, and the
stubgen docs are updated to the new form.
---
docs/packaging/stubgen.rst | 52 ++--
.../python/my_ffi_extension/_ffi_api.py | 6 +-
examples/rust_stubgen/README.md | 48 ++--
.../rust/src/generated/rust_stubgen/mod.rs | 9 +-
python/tvm_ffi/access_path.py | 4 +-
python/tvm_ffi/stub/cli.py | 94 ++++----
python/tvm_ffi/stub/consts.py | 30 +--
python/tvm_ffi/stub/directives.py | 113 +++++++++
python/tvm_ffi/stub/file_utils.py | 6 +-
python/tvm_ffi/stub/generator.py | 26 +-
python/tvm_ffi/stub/python_generator/codegen.py | 4 +-
python/tvm_ffi/stub/python_generator/consts.py | 7 +
python/tvm_ffi/stub/python_generator/generator.py | 26 +-
python/tvm_ffi/stub/python_generator/utils.py | 2 +-
python/tvm_ffi/stub/rust_generator/codegen.py | 34 +--
python/tvm_ffi/stub/rust_generator/consts.py | 18 +-
python/tvm_ffi/stub/rust_generator/directives.py | 217 ++++++++---------
python/tvm_ffi/stub/rust_generator/generator.py | 20 +-
python/tvm_ffi/stub/rust_generator/utils.py | 4 +-
python/tvm_ffi/testing/testing.py | 26 +-
tests/python/test_stubgen.py | 85 ++++---
tests/python/test_stubgen_rust.py | 268 +++++++++++----------
22 files changed, 592 insertions(+), 507 deletions(-)
diff --git a/docs/packaging/stubgen.rst b/docs/packaging/stubgen.rst
index f9271f28..ecf01dfe 100644
--- a/docs/packaging/stubgen.rst
+++ b/docs/packaging/stubgen.rst
@@ -384,45 +384,51 @@ When you run the tool, it:
# tvm-ffi-stubgen(end)
]
-``ty-map`` - Type Mapping
+``skip-file`` - Skip File
+ Prevents the tool from modifying the file. Place anywhere in the file.
+
+ .. code-block:: python
+
+ # tvm-ffi-stubgen(skip-file)
+
+**One-line directives**
+
+The other directives are settings on a single line, written
+``<subject>: <target>(<key>, <key> = <value>, ...)``; the parentheses are
omitted when
+there are no keys. Several lines for the same subject and target merge.
+
+``object: <type_key>(ty_map = <path>)`` - Type Mapping
Maps a C++ type key to a different Python type path. This is useful when
the Python
class is defined in a different module than the type key suggests.
.. code-block:: python
- # tvm-ffi-stubgen(ty-map): ffi.reflection.AccessStep ->
ffi.access_path.AccessStep
+ # tvm-ffi-stubgen(object): ffi.reflection.AccessStep(ty_map =
ffi.access_path.AccessStep)
-``prefix`` - Demand a Namespace
+``file: mod(prefix = <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, and a key the ``tvm_ffi`` crate binds itself (every
``ffi.*`` key)
- is never added. A block under the prefix whose type is no longer registered
is
- removed. Code outside the blocks is preserved.
+ registered directly under the prefix that no processed file defines yet; an
object's
+ ``skip`` key leaves one out, and a key the ``tvm_ffi`` crate binds itself
(every
+ ``ffi.*`` key) is never added. A block under the prefix whose type is no
longer
+ registered is removed. Code outside the blocks is preserved.
.. code-block:: rust
- // tvm-ffi-stubgen(prefix): my_ffi_extension
- // tvm-ffi-stubgen(skip): my_ffi_extension.Internal
+ // tvm-ffi-stubgen(file): mod(prefix = my_ffi_extension)
+ // tvm-ffi-stubgen(object): my_ffi_extension.Internal(skip)
A generated type is named from the module of the file its block is in, not
from its
type key: the module follows the file's place under the ``src/`` of the
nearest
``Cargo.toml`` (``src/ir/prim.rs`` is ``crate::ir::prim``). Other files
refer to it by
- that ``crate::`` path, so a namespace can live anywhere in the module tree.
+ that ``crate::`` path, so a namespace can live anywhere in the module tree.
The Rust
+ target's other keys are listed in ``examples/rust_stubgen``.
-``import-object`` - Import Object
- Injects a custom import into generated code. The format is
- ``<full_name>;<type_checking_only>;<alias>``.
+``import: <full_name>(type_checking_only, alias = <alias>)`` - Import Object
+ Keeps an import in the generated import section; both keys are optional.
.. code-block:: python
- # tvm-ffi-stubgen(import-object): ffi.Object;False;_ffi_Object
-
- This imports ``ffi.Object`` as ``_ffi_Object`` for use in generated code.
The second
- field (``False``) indicates the import is not TYPE_CHECKING-only.
-
-``skip-file`` - Skip File
- Prevents the tool from modifying the file. Place anywhere in the file.
-
- .. code-block:: python
+ # tvm-ffi-stubgen(import): ffi.Object(alias = _ffi_Object)
- # tvm-ffi-stubgen(skip-file)
+ This imports ``ffi.Object`` as ``_ffi_Object`` for use in generated code;
+ ``type_checking_only`` would import it under ``TYPE_CHECKING`` only.
diff --git a/examples/python_packaging/python/my_ffi_extension/_ffi_api.py
b/examples/python_packaging/python/my_ffi_extension/_ffi_api.py
index 0b67bd8b..de71f937 100644
--- a/examples/python_packaging/python/my_ffi_extension/_ffi_api.py
+++ b/examples/python_packaging/python/my_ffi_extension/_ffi_api.py
@@ -27,7 +27,7 @@ if TYPE_CHECKING:
# isort: on
# fmt: on
# tvm-ffi-stubgen(end)
-# tvm-ffi-stubgen(import-object):
tvm_ffi.libinfo.load_lib_module;False;_FFI_LOAD_LIB
+# tvm-ffi-stubgen(import): tvm_ffi.libinfo.load_lib_module(alias =
_FFI_LOAD_LIB)
LIB = _FFI_LOAD_LIB("my_ffi_extension", "my_ffi_extension")
# tvm-ffi-stubgen(begin): global/my_ffi_extension
# fmt: off
@@ -37,8 +37,8 @@ if TYPE_CHECKING:
def raise_error(_0: str, /) -> None: ...
# fmt: on
# tvm-ffi-stubgen(end)
-# tvm-ffi-stubgen(import-object): tvm_ffi.register_object;False;_FFI_REG_OBJ
-# tvm-ffi-stubgen(import-object): ffi.Object;False;_ffi_Object
+# tvm-ffi-stubgen(import): tvm_ffi.register_object(alias = _FFI_REG_OBJ)
+# tvm-ffi-stubgen(import): ffi.Object(alias = _ffi_Object)
@_FFI_REG_OBJ("my_ffi_extension.IntPair")
class IntPair(_ffi_Object):
"""FFI binding for `my_ffi_extension.IntPair`."""
diff --git a/examples/rust_stubgen/README.md b/examples/rust_stubgen/README.md
index d7ebeda0..60ac7c37 100644
--- a/examples/rust_stubgen/README.md
+++ b/examples/rust_stubgen/README.md
@@ -70,12 +70,12 @@ The Rust crate depends on the `tvm-ffi` crate of this
repository and needs
The binding file opens with a `prefix` line:
```rust
-// tvm-ffi-stubgen(prefix): rust_stubgen
+// tvm-ffi-stubgen(file): mod(prefix = rust_stubgen)
```
It makes the file own the `rust_stubgen` namespace: every object registered
directly under it gets an `object/<type_key>` block in this file on the next
-run, and a `skip` line (`// tvm-ffi-stubgen(skip): rust_stubgen.Internal`)
+run, and a `skip` key (`// tvm-ffi-stubgen(object):
rust_stubgen.Internal(skip)`)
leaves one out. Add a second object to `int_pair.cc` and rebuild, and its block
appears after `IntPair`; delete a block and it comes back, so dropping a
binding
means writing `skip`. Remove the object from the C++ side and its block is
@@ -87,12 +87,12 @@ file the command line points at. This example keeps them
under
## Hand-written constructors
-`IntPair::new` is hand-written. The `custom-new` directive makes the generator
+`IntPair::new` is hand-written. The object's `custom_new` key makes the
generator
name its lossless allocator `from_fields` instead of `new`, and the
hand-written constructor calls it after deriving `kind` from the operands:
```rust
-// tvm-ffi-stubgen(custom-new): rust_stubgen.IntPair
+// tvm-ffi-stubgen(object): rust_stubgen.IntPair(custom_new)
impl IntPair {
pub fn new(a: i64, b: i64) -> Self {
@@ -102,37 +102,39 @@ impl IntPair {
}
```
-Without the directive the block would define `IntPair::new` too and the crate
+Without the key the block would define `IntPair::new` too and the crate
would not compile.
## Directives
-The generated file keeps one-line directives the generator reads on every run.
-Besides `prefix` and `custom-new`, this example declares the integer field
-`kind` as an open newtype:
+The generated file keeps one-line directives the generator reads on every run,
+each `<subject>: <target>(<key>, <key> = <value>, ...)` (no parentheses without
+keys). Lines for the same subject and target merge. Besides `prefix` and
+`custom_new`, this example declares the integer field `kind` as an open
+newtype:
```rust
-// tvm-ffi-stubgen(enum): rust_stubgen.IntPair.kind -> PairKind(i32) {
Unordered=0, Ordered=1 }
+// tvm-ffi-stubgen(enum): PairKind(repr = i32, Unordered = 0, Ordered = 1)
+// tvm-ffi-stubgen(field): rust_stubgen.IntPair.kind(type = PairKind)
```
-Four more are available: `field` names the Rust type of a field
-(`// tvm-ffi-stubgen(field): rust_stubgen.IntPair.a -> MyInt`), `nullable`
-wraps it in `Option` (`// tvm-ffi-stubgen(nullable): rust_stubgen.IntPair.a`),
-`visibility` makes it `pub(crate)` instead of `pub`
-(`// tvm-ffi-stubgen(visibility): rust_stubgen.IntPair.a -> pub(crate)`), and
`upcast` adds a conversion to a hand-written typed view
-(`// tvm-ffi-stubgen(upcast): rust_stubgen.IntPair -> MyView`).
+A field also takes `type` naming any Rust type
+(`// tvm-ffi-stubgen(field): rust_stubgen.IntPair.a(type = MyInt)`), `nullable`
+to wrap it in `Option`, and `visibility = pub(crate)` instead of `pub`. An
+object takes `upcast = MyView` (repeatable) to add a conversion to a
+hand-written typed view, `opaque` to keep its layout opaque, and `no_new` to
+skip the generated allocator.
-When the C++ side fixes the value set in its ABI, `closed` before the member
-list makes `enum` emit a `#[repr(i32)]` Rust enum instead. Reading a value
-outside the set is undefined behavior, so the set should include 0, the value
of
-a field never set:
+When the C++ side fixes the value set in its ABI, `closed` makes `enum` emit a
+`#[repr(i32)]` Rust enum instead. Reading a value outside the set is undefined
+behavior, so the set should include 0, the value of a field never set:
```rust
-// tvm-ffi-stubgen(enum): rust_stubgen.IntPair.kind -> PairKind(i32) closed {
Unordered=0, Ordered=1 }
+// tvm-ffi-stubgen(enum): PairKind(repr = i32, closed, Unordered = 0, Ordered
= 1)
```
A crate built with `#![warn(missing_docs)]` adds
-`// tvm-ffi-stubgen(missing-docs): allow` to a file: every generated type and
+`// tvm-ffi-stubgen(file): mod(missing_docs = allow)` to a file: every
generated type and
`impl` in it gets `#[allow(missing_docs)]`, and docs written above a block
still document its first item.
@@ -142,9 +144,9 @@ Every type an object refers to (its parent, its ancestors,
the types of its
fields) has to be provided in the same `tvm-ffi-stubgen` run: `ffi.*` types
come from the `tvm-ffi` crate (a field of an `ffi.*` type the crate does not
bind is read as `Any`), an `object/<key>` block in any processed file
-generates it, and a `ty-map` points at a hand-written binding whose object
+generates it, and a `ty_map` key points at a hand-written binding whose object
struct is named `<Name>Obj`
-(`// tvm-ffi-stubgen(ty-map): rust_stubgen.IntPair -> crate::hand::IntPair`).
+(`// tvm-ffi-stubgen(object): rust_stubgen.IntPair(ty_map =
crate::hand::IntPair)`).
Anything else is an error listing the missing keys, so a partial binding never
references a module that does not exist. A hand-written binding may sit next to
the generated ones: the file of its own module refers to it by name instead of
diff --git a/examples/rust_stubgen/rust/src/generated/rust_stubgen/mod.rs
b/examples/rust_stubgen/rust/src/generated/rust_stubgen/mod.rs
index 1c7e1ad2..36d6d5c2 100644
--- a/examples/rust_stubgen/rust/src/generated/rust_stubgen/mod.rs
+++ b/examples/rust_stubgen/rust/src/generated/rust_stubgen/mod.rs
@@ -21,7 +21,7 @@
//! FFI bindings for `rust_stubgen` (generated by tvm-ffi-stubgen).
// Every object registered under `rust_stubgen` gets its block in this file;
`skip` leaves one out.
-// tvm-ffi-stubgen(prefix): rust_stubgen
+// tvm-ffi-stubgen(file): mod(prefix = rust_stubgen)
// tvm-ffi-stubgen(begin): import-section
use tvm_ffi::Error;
@@ -31,11 +31,12 @@ use tvm_ffi::Result;
use tvm_ffi::VALUE_ERROR;
// tvm-ffi-stubgen(end)
-// The `kind` field is an integer on the C++ side; this directive types it as
an
+// The `kind` field is an integer on the C++ side; these directives type it as
an
// open integer newtype in Rust.
-// tvm-ffi-stubgen(enum): rust_stubgen.IntPair.kind -> PairKind(i32) {
Unordered=0, Ordered=1 }
+// tvm-ffi-stubgen(enum): PairKind(repr = i32, Unordered = 0, Ordered = 1)
+// tvm-ffi-stubgen(field): rust_stubgen.IntPair.kind(type = PairKind)
// `IntPair::new` is hand-written below, so the generated allocator is
`from_fields`.
-// tvm-ffi-stubgen(custom-new): rust_stubgen.IntPair
+// tvm-ffi-stubgen(object): rust_stubgen.IntPair(custom_new)
// tvm-ffi-stubgen(begin): object/rust_stubgen.IntPair
#[repr(C)]
diff --git a/python/tvm_ffi/access_path.py b/python/tvm_ffi/access_path.py
index 8e887639..d690ba22 100644
--- a/python/tvm_ffi/access_path.py
+++ b/python/tvm_ffi/access_path.py
@@ -50,7 +50,7 @@ class AccessKind(IntEnum):
class AccessStep(Object):
"""Access step container."""
- # tvm-ffi-stubgen(ty-map): ffi.reflection.AccessStep ->
ffi.access_path.AccessStep
+ # tvm-ffi-stubgen(object): ffi.reflection.AccessStep(ty_map =
ffi.access_path.AccessStep)
# tvm-ffi-stubgen(begin): object/ffi.reflection.AccessStep
# fmt: off
kind: int
@@ -81,7 +81,7 @@ class AccessPath(Object):
"""
- # tvm-ffi-stubgen(ty-map): ffi.reflection.AccessPath ->
ffi.access_path.AccessPath
+ # tvm-ffi-stubgen(object): ffi.reflection.AccessPath(ty_map =
ffi.access_path.AccessPath)
# tvm-ffi-stubgen(begin): object/ffi.reflection.AccessPath
# fmt: off
parent: Object | None
diff --git a/python/tvm_ffi/stub/cli.py b/python/tvm_ffi/stub/cli.py
index eb3381c6..5a824273 100644
--- a/python/tvm_ffi/stub/cli.py
+++ b/python/tvm_ffi/stub/cli.py
@@ -29,6 +29,7 @@ from typing import TYPE_CHECKING
from tvm_ffi.core import _object_type_key_to_index
from . import consts as C
+from .directives import Directives, Value
from .file_utils import CodeBlock, FileInfo, collect_files, syntax_for
from .generator import generator_names, get_generator
from .layout import classify, write_coverage_report
@@ -70,20 +71,14 @@ def __main__() -> int:
init_path = init_path.parent
# Stage 1: Collect information
- # - type maps: `tvm-ffi-stubgen(ty-map)`
+ # - type maps: `tvm-ffi-stubgen(object): <type_key>(ty_map = <name>)`
# - defined global functions: `tvm-ffi-stubgen(begin): global/...`
# - defined object types: `tvm-ffi-stubgen(begin): object/...`
ty_map: dict[str, str] = generator.default_ty_map()
failed = 0
stale = 0
for file in files:
- try:
- _stage_1(file, ty_map)
- except Exception:
- failed += 1
- print(
- f'{C.TERM_RED}[Failed] File "{file.path}":
{traceback.format_exc()}{C.TERM_RESET}'
- )
+ _stage_1(file, ty_map, generator)
# Stage 2. Generate stubs if they are not defined on the file.
generated_prefixes: set[str] = set()
@@ -101,8 +96,8 @@ def __main__() -> int:
# `--check` compares the generator's view of each file before the run
edits any block.
before = {file.path: generator.check_view(file.code_blocks) for file in
files if opt.check}
- # 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.
+ # Stage 2b. Add the object blocks a `tvm-ffi-stubgen(file): mod(prefix =
...)` file asks
+ # for. This runs after `--init`, which rewrites files on disk and reloads
them.
failed += _roll_out_prefixes(files, generator)
# Stage 3: Process
@@ -150,20 +145,31 @@ def __main__() -> int:
return 2 if failed else 1 if stale else 0
-def _stage_1(
- file: FileInfo,
- ty_map: dict[str, str],
-) -> None:
+def _directives(file: FileInfo, generator: Generator) -> Directives:
+ """Parse the one-line directives of ``file`` against the pipeline's and
the generator's keys."""
+ pipeline, own = C.PIPELINE_DIRECTIVE_KEYS, generator.directive_keys
+ directives = Directives(
+ {s: {**pipeline.get(s, {}), **own.get(s, {})} for s in {*pipeline,
*own}}
+ )
for code in file.code_blocks:
- 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()
+ if code.kind == "directive":
+ subject, payload = code.param
+ directives.add(subject, payload, code.lineno_start)
+ return directives
+
+
+def _settings(file: FileInfo, generator: Generator) -> dict[str, dict[str,
dict[str, Value]]]:
+ """Return the settings the pipeline reads ahead of ``_stage_3``, which
reports a bad file."""
+ try:
+ return _directives(file, generator).settings
+ except ValueError:
+ return {}
+
+
+def _stage_1(file: FileInfo, ty_map: dict[str, str], generator: Generator) ->
None:
+ for type_key, keys in _settings(file, generator).get("object", {}).items():
+ if "ty_map" in keys:
+ ty_map[type_key] = str(keys["ty_map"])
def _roll_out_prefixes(files: list[FileInfo], generator: Generator) -> int:
@@ -179,20 +185,17 @@ def _roll_out_prefixes(files: list[FileInfo], generator:
Generator) -> int:
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:
+ settings = _settings(file, generator)
+ if "prefix" not in settings.get("file", {}).get("mod", {}):
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:
+ prefix = str(settings["file"]["mod"]["prefix"]).rstrip(".")
+ head = next(c for c in file.code_blocks if c.kind == "directive" and
c.param[0] == "file")
+ if owners.setdefault(prefix, file.path) != file.path:
failed += 1
- print(f'{C.TERM_RED}[Failed] File "{file.path}":
{error}{C.TERM_RESET}')
+ print(
+ f'{C.TERM_RED}[Failed] File "{file.path}": prefix `{prefix}`
is already '
+ f"declared by {owners[prefix]}{C.TERM_RESET}"
+ )
continue
_remove_unregistered_blocks(file, prefix, registry.get(prefix, []))
if prefix not in registry:
@@ -200,7 +203,7 @@ def _roll_out_prefixes(files: list[FileInfo], generator:
Generator) -> int:
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"}
+ skipped = {key for key, keys in settings.get("object", {}).items() if
"skip" in keys}
keys = [
key
for key in registry[prefix]
@@ -279,10 +282,10 @@ def _stage_2(
code.param for file in files for code in file.code_blocks if code.kind
== "object"
}
skipped: set[str] = {
- code.param[1].strip()
+ type_key
for file in files
- for code in file.code_blocks
- if code.kind == "directive" and code.param[0] == "skip"
+ for type_key, keys in _settings(file, generator).get("object",
{}).items()
+ if "skip" in keys
}
# Step 0. Generate missing `_ffi_api.py` and `__init__.py` under each
prefix.
@@ -367,17 +370,8 @@ def _stage_3( # noqa: PLR0912
defined_funcs: set[str] = set()
unregistered: list[str] = []
defined_types: set[str] = set()
- imports = generator.new_imports()
- # 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 != "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 1. Parse the one-line directives and seed the generator's
collector with them.
+ imports = generator.new_imports(_directives(file, generator))
# 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 5cb36304..6b4dcfb0 100644
--- a/python/tvm_ffi/stub/consts.py
+++ b/python/tvm_ffi/stub/consts.py
@@ -23,6 +23,8 @@ from typing import Literal
from typing_extensions import TypeAlias
+from .directives import Key, KeyTable
+
@dataclasses.dataclass(frozen=True)
class MarkerSyntax:
@@ -52,24 +54,14 @@ class MarkerSyntax:
"""Marker that closes a generated block: ``<comment>
tvm-ffi-stubgen(end)``."""
return f"{self.prefix}end)"
- @property
- def ty_map(self) -> str:
- """The ``ty-map`` directive marker: ``<comment>
tvm-ffi-stubgen(ty-map):``."""
- return self.directive("ty-map")
-
- @property
- def import_object(self) -> str:
- """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}):"
+ def directive(self, subject: str) -> str:
+ """One-line directive marker: ``<comment>
tvm-ffi-stubgen(<subject>):``."""
+ return f"{self.prefix}{subject}):"
PYTHON_SYNTAX = MarkerSyntax(comment="#")
@@ -84,10 +76,14 @@ 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", "prefix",
"skip"})
+#: Directive keys the language-neutral pipeline reads (see
:mod:`.directives`): an
+#: ``object`` maps its type key (``ty_map``) or opts out of ``--init`` and
``prefix``
+#: (``skip``); ``file: mod(prefix = <p>)`` asks for a block per object under
``p``.
+#: The active generator declares every other key
(``Generator.directive_keys``).
+PIPELINE_DIRECTIVE_KEYS: KeyTable = {
+ "object": {"skip": Key(), "ty_map": Key(takes_value=True)},
+ "file": {"prefix": Key(takes_value=True)},
+}
STUB_BLOCK_KINDS: TypeAlias = Literal[
"global",
diff --git a/python/tvm_ffi/stub/directives.py
b/python/tvm_ffi/stub/directives.py
new file mode 100644
index 00000000..2d362e0a
--- /dev/null
+++ b/python/tvm_ffi/stub/directives.py
@@ -0,0 +1,113 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""The one-line directive grammar shared by every generator.
+
+A directive names a subject, a target, and optional keys::
+
+ <comment> tvm-ffi-stubgen(<subject>): <target>
+ <comment> tvm-ffi-stubgen(<subject>): <target>(<key>, <key> = <value>, ...)
+
+A key is a flag or takes a value; commas and ``=`` inside ``()``, ``[]``,
``{}`` or
+``<>`` belong to the value, so ``pub(crate)`` and ``Option<A, B>`` are one
value.
+Lines for the same subject and target merge: a repeatable key collects every
value,
+any other key may repeat only with the same value. The pipeline reads the
+``object`` / ``file`` keys of :data:`consts.PIPELINE_DIRECTIVE_KEYS`; the
active
+generator declares the rest (``Generator.directive_keys``).
+"""
+
+from __future__ import annotations
+
+import dataclasses
+from collections.abc import Mapping
+from typing import Union
+
+#: A parsed value: ``True`` for a flag, the text, or every text of a
repeatable key.
+Value = Union[bool, str, list[str]]
+
+
[email protected](frozen=True)
+class Key:
+ """How a key is written: a flag, or ``key = value`` (``values`` lists the
legal ones)."""
+
+ takes_value: bool = False
+ repeatable: bool = False
+ values: tuple[str, ...] = ()
+
+
+#: Key tables: subject -> key -> spelling. ``"*"`` accepts any other key with
a value.
+KeyTable = Mapping[str, Mapping[str, Key]]
+
+
+class Directives:
+ """The directives of one file: ``settings[subject][target][key]``."""
+
+ def __init__(self, keys: KeyTable) -> None:
+ self.keys = keys
+ self.settings: dict[str, dict[str, dict[str, Value]]] = {s: {} for s
in keys}
+ self._lines: dict[tuple[str, str, str], int] = {}
+
+ def add(self, subject: str, payload: str, lineno: int) -> None:
+ """Parse and merge one directive; raise ``ValueError`` on a malformed
one."""
+ if subject not in self.keys:
+ raise ValueError(f"Unknown directive `{subject}` at line {lineno}")
+
+ def invalid(reason: str) -> ValueError:
+ return ValueError(f"Invalid `{subject}` directive at line
{lineno}: {reason}")
+
+ target, paren, rest = payload.strip().partition("(")
+ target = target.strip()
+ if paren and not rest.endswith(")"):
+ raise invalid("expected `<target>(<key>, <key> = <value>, ...)`")
+ if not target or any(c.isspace() for c in target):
+ raise invalid(f"bad target `{target}`")
+ if (subject == "field" and "." not in target.strip(".")) or (
+ subject == "file" and target != "mod"
+ ):
+ raise invalid("the target is `<type_key>.<field>`" if subject ==
"field" else "`mod`")
+ entry = self.settings[subject].setdefault(target, {})
+ for arg in filter(str.strip, _split(rest[:-1], ",")):
+ key, *value_parts = _split(arg, "=")
+ key, eq, value = key.strip(), bool(value_parts),
"=".join(value_parts).strip()
+ spec = self.keys[subject].get(key) or self.keys[subject].get("*")
+ if not key.isidentifier() or spec is None:
+ raise invalid(f"unknown key `{key}`")
+ if eq != spec.takes_value or (eq and not value):
+ raise invalid(
+ f"`{key}` takes a value" if spec.takes_value else
f"`{key}` is a flag"
+ )
+ if spec.values and value not in spec.values:
+ raise invalid(f"`{key}` is one of {', '.join(spec.values)}")
+ if spec.repeatable:
+ collected = entry.setdefault(key, [])
+ assert isinstance(collected, list)
+ collected.append(value)
+ continue
+ new: Value = value if eq else True
+ first = self._lines.setdefault((subject, target, key), lineno)
+ if entry.setdefault(key, new) != new:
+ raise invalid(f"`{key}` conflicts with line {first}")
+
+
+def _split(text: str, sep: str) -> list[str]:
+ """Split ``text`` on each ``sep`` outside brackets."""
+ parts, depth, start = [], 0, 0
+ for i, char in enumerate(text):
+ depth += (char in "([{<") - (char in ")]}>")
+ if char == sep and depth == 0:
+ parts.append(text[start:i])
+ start = i + 1
+ return [*parts, text[start:]]
diff --git a/python/tvm_ffi/stub/file_utils.py
b/python/tvm_ffi/stub/file_utils.py
index 381254cf..8e93ed42 100644
--- a/python/tvm_ffi/stub/file_utils.py
+++ b/python/tvm_ffi/stub/file_utils.py
@@ -68,8 +68,8 @@ class CodeBlock:
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 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.
+ # One-line directive `<comment> tvm-ffi-stubgen(<subject>):
<payload>`. The payload is
+ # kept verbatim for `directives.Directives` to parse.
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}")
@@ -188,7 +188,7 @@ class FileInfo:
codes.append(code)
code = None
elif clean_line.startswith(syntax.prefix):
- # Process a one-line directive "<comment>
tvm-ffi-stubgen(<name>): ..."
+ # Process a one-line directive "<comment>
tvm-ffi-stubgen(<subject>): ..."
dir_code = CodeBlock.from_begin_line(lineno, clean_line,
syntax)
dir_code.lineno_end = lineno
dir_code.lines.append(line)
diff --git a/python/tvm_ffi/stub/generator.py b/python/tvm_ffi/stub/generator.py
index 3a8c6885..1c6b3407 100644
--- a/python/tvm_ffi/stub/generator.py
+++ b/python/tvm_ffi/stub/generator.py
@@ -30,8 +30,8 @@ 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 the one-line directives the
-generator declares, and later render it, but never reaches inside. Adding a
+asks the generator to create one from the file's one-line directives
+(:mod:`.directives`), and later render it, but never reaches inside. Adding a
language is therefore "implement one more :class:`Generator`" rather than
forking the pipeline.
"""
@@ -48,6 +48,7 @@ if TYPE_CHECKING:
from collections.abc import Mapping
from pathlib import Path
+ from .directives import Directives, KeyTable
from .file_utils import CodeBlock
from .utils import FuncInfo, InitConfig, ObjectInfo, Options
@@ -77,10 +78,10 @@ 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]
+ #: Keys of the one-line directives (``<comment>
tvm-ffi-stubgen(<subject>): <target>(...)``)
+ #: this generator reads, per subject; merged with
:data:`consts.PIPELINE_DIRECTIVE_KEYS`.
+ #: A subject or key in neither table is an error.
+ directive_keys: KeyTable
def default_ty_map(self) -> dict[str, str]:
"""Return the default FFI-origin -> target-type name map for this
language."""
@@ -88,16 +89,11 @@ class Generator(Protocol):
# --- import collection (representation is generator-private) ------------
- def new_imports(self) -> Any:
- """Create a fresh, empty import collector for one file."""
- ...
-
- def add_directive(self, imports: Any, name: str, payload: str, lineno:
int) -> None:
- """Record a one-line directive (raw payload) into ``imports``.
+ def new_imports(self, directives: Directives) -> Any:
+ """Create the import collector for one file, seeded from the file's
``directives``.
- 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.
+ The collector is per file, so a directive applies to the blocks of the
file it
+ appears in.
"""
...
diff --git a/python/tvm_ffi/stub/python_generator/codegen.py
b/python/tvm_ffi/stub/python_generator/codegen.py
index 2669231d..40b9c53f 100644
--- a/python/tvm_ffi/stub/python_generator/codegen.py
+++ b/python/tvm_ffi/stub/python_generator/codegen.py
@@ -64,7 +64,7 @@ class {type_name}({parent_type_name}):
def _prompt_import_object(type_key: str, type_name: str, syntax:
C.MarkerSyntax) -> str:
- return f"""{syntax.import_object} {type_key};False;{type_name}\n"""
+ return f"""{syntax.directive("import")} {type_key}(alias =
{type_name})\n"""
def _prompt_import_section(syntax: C.MarkerSyntax) -> str:
@@ -105,7 +105,7 @@ def _type_suffix_and_record(
def _make_input_ty_map(ty_map: dict[str, str]) -> dict[str, str]:
- """Derive input-side defaults without overriding explicit ty-map
entries."""
+ """Derive input-side defaults without overriding explicit ty_map
entries."""
input_ty_map = ty_map.copy()
for key, input_default in PC.TY_MAP_INPUT_DEFAULTS.items():
output_default = PC.TY_MAP_DEFAULTS.get(key)
diff --git a/python/tvm_ffi/stub/python_generator/consts.py
b/python/tvm_ffi/stub/python_generator/consts.py
index a2bee13d..8dfbf68f 100644
--- a/python/tvm_ffi/stub/python_generator/consts.py
+++ b/python/tvm_ffi/stub/python_generator/consts.py
@@ -24,6 +24,13 @@ typing assumptions.
from __future__ import annotations
+from ..directives import Key, KeyTable
+
+#: ``import: a.b.C(type_checking_only, alias = _C)`` keeps ``a.b.C`` in the
import section.
+PYTHON_DIRECTIVE_KEYS: KeyTable = {
+ "import": {"type_checking_only": Key(), "alias": Key(takes_value=True)},
+}
+
#: Default FFI-origin -> Python-type name map used to seed a render.
TY_MAP_DEFAULTS = {
"Any": "typing.Any",
diff --git a/python/tvm_ffi/stub/python_generator/generator.py
b/python/tvm_ffi/stub/python_generator/generator.py
index 40a47bfa..231eb5ac 100644
--- a/python/tvm_ffi/stub/python_generator/generator.py
+++ b/python/tvm_ffi/stub/python_generator/generator.py
@@ -35,6 +35,7 @@ if TYPE_CHECKING:
from collections.abc import Mapping
from pathlib import Path
+ from ..directives import Directives
from ..file_utils import CodeBlock
from ..utils import FuncInfo, InitConfig, ObjectInfo, Options
@@ -45,7 +46,7 @@ class PythonGenerator:
name = "python"
syntax = C.PYTHON_SYNTAX
source_exts = frozenset({".py", ".pyi"})
- directive_kinds: frozenset[str] = frozenset({"import-object"})
+ directive_keys = PC.PYTHON_DIRECTIVE_KEYS
def default_ty_map(self) -> dict[str, str]:
"""Return the default FFI-origin -> Python-type name map."""
@@ -53,19 +54,16 @@ class PythonGenerator:
# --- import collection (Python representation is private) ---------------
- def new_imports(self) -> PythonImports:
- """Create an empty import collector."""
- return PythonImports()
-
- 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(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 new_imports(self, directives: Directives) -> PythonImports:
+ """Create an import collector holding the file's ``import``
directives."""
+ imports = PythonImports()
+ for full_name, keys in directives.settings["import"].items():
+ alias = str(keys["alias"]) if "alias" in keys else None
+ tco = "type_checking_only" in keys
+ imports.items.append(ImportItem(full_name, type_checking_only=tco,
alias=alias))
+ if alias == "_FFI_LOAD_LIB" or
full_name.endswith("libinfo.load_lib_module"):
+ imports.has_lib_load = True
+ return imports
def canonical_type_name(self, type_key: str) -> str:
"""Return the canonical (import-comparable) full name for a defined
type key."""
diff --git a/python/tvm_ffi/stub/python_generator/utils.py
b/python/tvm_ffi/stub/python_generator/utils.py
index e4d6dfc9..bfecc02d 100644
--- a/python/tvm_ffi/stub/python_generator/utils.py
+++ b/python/tvm_ffi/stub/python_generator/utils.py
@@ -89,7 +89,7 @@ class PythonImports:
"""Opaque import collector threaded through the Python generation pipeline.
The language-agnostic ``cli`` treats this as an opaque token: it asks the
- generator to create one, seed it from ``import-object`` directives, and
later
+ generator to create one, seed it from ``import`` directives, and later
render it. Only the Python generator reaches inside.
"""
diff --git a/python/tvm_ffi/stub/rust_generator/codegen.py
b/python/tvm_ffi/stub/rust_generator/codegen.py
index 708dae3e..858650e8 100644
--- a/python/tvm_ffi/stub/rust_generator/codegen.py
+++ b/python/tvm_ffi/stub/rust_generator/codegen.py
@@ -25,22 +25,22 @@ the verdict of :mod:`tvm_ffi.stub.layout`:
width, pinned by a ``const`` size/alignment assertion. ``<Leaf>Obj::new``
and the wrapper's ``new`` take every field root to leaf: the former each as
its own type, the latter a scalar as its own type, an object reference as
- ``impl IntoArg`` and anything else as ``impl Into``. ``custom-new``
+ ``impl IntoArg`` and anything else as ``impl Into``. ``custom_new``
leaves the wrapper's ``new`` to hand-written code and names the generated
one ``from_fields``. A type that is not final
can be the parent of a type defined in Rust (``tvm_ffi::rs_object``): its
``<Leaf>Obj::new`` is public, and its object derives ``ShallowCopy`` and
``UnsafeInit`` (``rs_parent``) when every ancestor's is generated
- too (a ``ty-map``'d one is hand-written and may not). ``<Leaf>Obj::new`` of
a final type is crate-private.
+ too (a ``ty_map``'d one is hand-written and may not). ``<Leaf>Obj::new`` of
a final type is crate-private.
- *opaque*: the struct embeds only its parent, and each field is read through
the C ABI getter. Nothing allocates it.
A field without a Rust mirror (``Optional<Any>``, ``Union``, ``void*``, an
-``ffi.*`` type the default ``ty-map`` does not list, a C++ standard-library
type
+``ffi.*`` type the default ``ty_map`` does not list, a C++ standard-library
type
such as ``std::optional``, an integer narrower than 8 bytes, whose signedness
the
schema does not record, or any field whose reflected size or alignment is not
its
-mirror's) makes the type opaque and is read through its getter; a ``field`` or
-``enum`` directive types a narrow integer; an ``opaque`` directive vetoes a
+mirror's) makes the type opaque and is read through its getter; a field's
``type``
+(a scalar or an ``enum``) types a narrow integer; an object's ``opaque``
vetoes a
reproducible layout; a scalar width named by a directive is checked against the
reflected field size.
A builtin parent the crate defines (``ffi.Enum``:
``tvm_ffi::enum_ref::EnumObj``) is
@@ -52,7 +52,7 @@ An enum type (under ``ffi.Enum``) has no allocator, since its
registry
allocates the values, and its reference implements ``tvm_ffi::EnumRef``.
Every type key an object refers to (parent, ancestors, field types) must be
-provided in the same run: ``ffi.*`` by the crate, a ``ty-map`` by a
hand-written
+provided in the same run: ``ffi.*`` by the crate, a ``ty_map`` by a
hand-written
binding whose object struct is ``<Name>Obj``, anything else by an ``object/``
block in one of the processed files. Otherwise the block is an error naming the
missing keys, so a partial binding never references a module that does not
exist.
@@ -165,7 +165,7 @@ def _is_object_ref(schema: TypeSchema) -> bool:
"""Whether the Rust type of ``schema`` is an object reference.
A builtin origin is one when listed in ``RUST_OBJECT_REF_ORIGINS``; a type
key
- outside ``ffi.*`` always is, as a generated or ``ty-map``'d binding.
+ outside ``ffi.*`` always is, as a generated or ``ty_map``'d binding.
"""
origin = schema.origin
if origin in C_RUST.RUST_OBJECT_REF_ORIGINS:
@@ -258,7 +258,7 @@ class _ObjectRenderer:
) -> str | None:
"""Resolve a leaf origin to its in-scope Rust name (recording its
``use``), or ``None``.
- A ``ty-map`` entry for the type key the schema was parsed from wins
over the origin.
+ A ``ty_map`` entry for the type key the schema was parsed from wins
over the origin.
"""
key = type_key if type_key in self.ty_map else origin
mapped = self.ty_map.get(key)
@@ -266,12 +266,12 @@ class _ObjectRenderer:
if "." not in origin or origin.startswith("ctypes."):
return None
if self._provider(origin) == "crate": # a key nobody provides is
recorded
- return None # an `ffi.*` key off the default `ty-map`: no
crate type, read as `Any`
+ return None # an `ffi.*` key off the default `ty_map`: no
crate type, read as `Any`
return imports.record(self._generated_type_path(origin))
return self._record_mapped(key, mapped, imports)
def _record_mapped(self, type_key: str, path: str, imports: RustImports)
-> str:
- """Record the ``ty-map`` target of ``type_key``; one in this file's
own module is in scope.
+ """Record the ``ty_map`` target of ``type_key``; one in this file's
own module is in scope.
A hand-written binding next to the generated ones (``ir.Op ->
crate::ir::Op``
in ``crate::ir``) must not import itself.
@@ -303,7 +303,7 @@ class _ObjectRenderer:
def _provider(self, type_key: str) -> str | None:
"""Who provides the binding of ``type_key``; a key nobody does is
recorded in :attr:`missing`.
- ``"crate"`` for builtin ``ffi.*`` types, ``"mapped"`` for a ``ty-map``
to a
+ ``"crate"`` for builtin ``ffi.*`` types, ``"mapped"`` for a ``ty_map``
to a
hand-written binding, ``"generated"`` for an ``object/`` block in this
run.
"""
if is_crate_type_key(type_key):
@@ -316,7 +316,7 @@ class _ObjectRenderer:
return None
def _base_type(self) -> tuple[str, bool]:
- """Resolve the ``base_`` struct and whether it is a generated or
``ty-map``'d parent.
+ """Resolve the ``base_`` struct and whether it is a generated or
``ty_map``'d parent.
A builtin parent below ``ffi.Object`` is embedded as the crate's object
struct, or as its header-only stand-in over the deepest builtin
ancestor the
@@ -441,7 +441,7 @@ class _ObjectRenderer:
# --- pieces ------------------------------------------------------------
def _allow_missing_docs(self) -> list[str]:
- """``#[allow(missing_docs)]`` under ``missing-docs: allow``, for a
crate that warns on it."""
+ """``#[allow(missing_docs)]`` under ``missing_docs = allow``, for a
crate that warns on it."""
return ["#[allow(missing_docs)]"] if
self.imports.directives.allow_missing_docs else []
def _visibility(self, target: str) -> str:
@@ -566,7 +566,7 @@ class _ObjectRenderer:
"""``impl_object_ref_upcast!`` to every ancestor's wrapper, then the
``upcast`` directives.
A generated ancestor accepts every object of its type, so its upcast
takes the
- compile-time checked form; a ``ty-map``'d one or an ``upcast`` view
may promise
+ compile-time checked form; a ``ty_map``'d one or an ``upcast`` view
may promise
more, so it takes the ``unsafe`` form.
"""
checked: list[str] = []
@@ -727,7 +727,7 @@ class _ObjectRenderer:
field as its own type, as ``#[rs_object]``'s ``new`` does, an object
reference as
``impl IntoArg`` and any other as ``impl Into`` (:func:`_param_type`),
and
converts it on the way to ``<Leaf>Obj::new``.
- ``custom-new`` names the wrapper's allocator ``from_fields`` instead;
``no-new``
+ ``custom_new`` names the wrapper's allocator ``from_fields`` instead;
``no_new``
skips it, and ``<Leaf>Obj::new`` of a final type with it.
``<Leaf>Obj::new`` of a type that is not final is public: a subtype
defined in Rust
builds its parent's object with it.
@@ -814,7 +814,7 @@ class _ObjectRenderer:
A subtype defined in Rust needs both of its parent's object. They are
derived
for a complete type that is not final and whose ancestors below
``ffi.Object``
are all generated: each of those derives them by this rule, while a
- ``ty-map``'d one is hand-written and may not.
+ ``ty_map``'d one is hand-written and may not.
"""
if self.info.is_final or not verdict.is_complete:
return False
@@ -920,7 +920,7 @@ def generate_rust_object(
keys = ", ".join(f"`{key}`" for key in sorted(renderer.missing))
raise ValueError(
f"`{obj_info.type_key}` (line {code.lineno_start}) depends on
{keys}, which this run "
- "does not ask for: add an `object/<key>` block or a `ty-map`
directive for each"
+ "does not ask for: add an `object/<key>` block or a `ty_map`
directive key for each"
)
indent = " " * code.indent
code.lines = [
diff --git a/python/tvm_ffi/stub/rust_generator/consts.py
b/python/tvm_ffi/stub/rust_generator/consts.py
index 01e1d521..091f4377 100644
--- a/python/tvm_ffi/stub/rust_generator/consts.py
+++ b/python/tvm_ffi/stub/rust_generator/consts.py
@@ -18,22 +18,6 @@
from __future__ import annotations
-#: One-line directives the Rust backend consumes.
-RUST_DIRECTIVE_KINDS = frozenset(
- {
- "import-object",
- "field",
- "nullable",
- "visibility",
- "enum",
- "opaque",
- "upcast",
- "custom-new",
- "no-new",
- "missing-docs",
- }
-)
-
#: Default FFI-origin -> Rust-type map; ``::`` paths get a ``use``, bare names
do not.
RUST_TY_MAP_DEFAULTS = {
"int": "i64",
@@ -120,7 +104,7 @@ RUST_EXACT_PARAMS = frozenset(
)
#: Builtin origins whose Rust type is an object reference, taken as ``impl
IntoArg<T>``.
-#: A type key outside ``ffi.*`` is one too: a generated or ``ty-map``'d
binding. ``Shape``,
+#: A type key outside ``ffi.*`` is one too: a generated or ``ty_map``'d
binding. ``Shape``,
#: like ``Array`` / ``List`` / ``Map``, is a container taken as ``impl
Into<T>``, which also
#: accepts what it is built from (``From<impl AsRef<[i64]>>``).
RUST_OBJECT_REF_ORIGINS = frozenset(
diff --git a/python/tvm_ffi/stub/rust_generator/directives.py
b/python/tvm_ffi/stub/rust_generator/directives.py
index 4ca394a4..0065c511 100644
--- a/python/tvm_ffi/stub/rust_generator/directives.py
+++ b/python/tvm_ffi/stub/rust_generator/directives.py
@@ -14,46 +14,63 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""The Rust backend's one-line directives: payload grammar and per-file
storage.
-
-Four address one reflected field as ``<type_key>.<field>``, three address a
type,
-and one applies to the whole file::
-
- // tvm-ffi-stubgen(field): tirx.Add.a -> PrimExpr
- // tvm-ffi-stubgen(nullable): ir.Expr.span
- // tvm-ffi-stubgen(visibility): ir.IRModule.global_var_map_ -> pub(crate)
- // tvm-ffi-stubgen(enum): tirx.For.kind -> ForKind(i32) { Serial=0,
Parallel=1 }
- // tvm-ffi-stubgen(opaque): ir.SourceName
- // tvm-ffi-stubgen(upcast): tirx.Add -> PrimExpr
- // tvm-ffi-stubgen(custom-new): tirx.Add
- // tvm-ffi-stubgen(no-new): ir.Expr
- // tvm-ffi-stubgen(missing-docs): allow
-
-``field`` sets the field's Rust type (a name in scope, or a ``::`` path to
-``use``); on a field inherited from an ancestor it narrows the allocator
-parameter instead. ``nullable`` wraps it in ``Option``; ``visibility`` makes
the
-field and its accessor ``pub(crate)`` instead of ``pub``; ``enum`` declares an
open
-integer newtype for it, or with ``closed`` before the member list a
-``#[repr(<int>)]`` enum of exactly those members; ``opaque`` keeps a type
opaque
-even when its layout is reproducible. ``upcast`` adds an upcast to a reference
-outside the ancestor chain; ``custom-new`` says the wrapper's ``new`` is
-hand-written; the generated one is named ``from_fields`` instead. ``no-new``
-skips the wrapper's generated allocator, and for a final type
``<Leaf>Obj::new``
-too, which only that allocator uses.
-``missing-docs: allow`` puts ``#[allow(missing_docs)]`` on the file's generated
-types and ``impl`` blocks.
+"""The Rust backend's directive keys, and the settings they give one file.
+
+Each subject's keys (see :mod:`tvm_ffi.stub.directives` for the grammar)::
+
+ // tvm-ffi-stubgen(object): tirx.Add(custom_new, upcast = PrimExpr)
+ // tvm-ffi-stubgen(object): ir.SourceName(opaque)
+ // tvm-ffi-stubgen(object): ir.Expr(no_new)
+ // tvm-ffi-stubgen(field): tirx.Add.a(type = PrimExpr)
+ // tvm-ffi-stubgen(field): ir.Expr.span(nullable)
+ // tvm-ffi-stubgen(field): ir.IRModule.global_var_map_(visibility =
pub(crate))
+ // tvm-ffi-stubgen(enum): ForKind(repr = i32, Serial = 0, Parallel = 1)
+ // tvm-ffi-stubgen(field): tirx.For.kind(type = ForKind)
+ // tvm-ffi-stubgen(file): mod(missing_docs = allow)
+ // tvm-ffi-stubgen(import): tvm_ffi::AnyView
+
+``object`` keys: ``opaque`` keeps a type opaque even when its layout is
+reproducible; ``upcast`` (repeatable) adds an upcast to a reference outside the
+ancestor chain; ``custom_new`` says the wrapper's ``new`` is hand-written, so
the
+generated one is named ``from_fields``; ``no_new`` skips the wrapper's
generated
+allocator, and for a final type ``<Leaf>Obj::new`` too, which only that
allocator
+uses. ``field`` keys: ``type`` sets the field's Rust type (a name in scope, a
+``::`` path to ``use``, or an ``enum`` of the file); on a field inherited from
an
+ancestor it narrows the allocator parameter instead. ``nullable`` wraps it in
+``Option``; ``visibility = pub(crate)`` makes the field and its accessor
+``pub(crate)`` instead of ``pub``. ``enum`` declares an open integer newtype of
+``repr`` with the members ``Name = <int>``, or with ``closed`` a
``#[repr(<int>)]``
+enum of exactly those members. ``file: mod(missing_docs = allow)`` puts
+``#[allow(missing_docs)]`` on the file's generated types and ``impl`` blocks.
+``import`` keeps a path ``use``d in the import section.
"""
from __future__ import annotations
import dataclasses
-import re
+from typing import TYPE_CHECKING
-_ENUM_RE = re.compile(
-
r"^(?P<target>\S+)\s*->\s*(?P<name>[A-Za-z_]\w*)\((?P<repr>[iu](?:8|16|32|64))\)"
- r"\s*(?:(?P<closed>closed\s*)?\{(?P<body>[^{}]*)\})?$"
-)
-_MEMBER_RE = re.compile(r"^(?P<name>[A-Za-z_]\w*)\s*=\s*(?P<value>-?\d+)$")
+from ..directives import Key, KeyTable
+
+if TYPE_CHECKING:
+ from ..directives import Directives
+
+_FLAG, _VALUE = Key(), Key(takes_value=True)
+_INT_TYPES = ("i8", "i16", "i32", "i64", "u8", "u16", "u32", "u64")
+
+#: The keys the Rust backend reads; an ``enum``'s other keys are its members.
+RUST_DIRECTIVE_KEYS: KeyTable = {
+ "object": {
+ "custom_new": _FLAG,
+ "no_new": _FLAG,
+ "opaque": _FLAG,
+ "upcast": Key(takes_value=True, repeatable=True),
+ },
+ "field": {"type": _VALUE, "nullable": _FLAG, "visibility": Key(True,
values=("pub(crate)",))},
+ "enum": {"repr": Key(True, values=_INT_TYPES), "closed": _FLAG, "*":
_VALUE},
+ "file": {"missing_docs": Key(True, values=("allow",))},
+ "import": {},
+}
@dataclasses.dataclass(frozen=True)
@@ -72,8 +89,8 @@ class EnumSpec:
@dataclasses.dataclass
-class Directives:
- """The Rust directives of one file, keyed by ``<type_key>.<field>`` or
``<type_key>``."""
+class RustDirectives:
+ """The Rust settings of one file, keyed by ``<type_key>.<field>`` or
``<type_key>``."""
field_types: dict[str, str] = dataclasses.field(default_factory=dict)
nullable: set[str] = dataclasses.field(default_factory=set)
@@ -85,91 +102,49 @@ class Directives:
no_new: set[str] = dataclasses.field(default_factory=set)
allow_missing_docs: bool = False
- def add(self, name: str, payload: str, lineno: int) -> None:
- """Parse and store one directive; raise ``ValueError`` on a malformed
payload."""
- if name == "field":
- lhs, rust_type = _split_arrow(name, payload, lineno,
"<type_key>.<field> -> <RustType>")
- self.field_types[_field_target(name, lhs, lineno)] = rust_type
- elif name == "nullable":
- self.nullable.add(_field_target(name, payload, lineno))
- elif name == "visibility":
- lhs, vis = _split_arrow(name, payload, lineno, "<type_key>.<field>
-> pub(crate)")
- if vis != "pub(crate)": # the one non-default level so far
- raise _invalid(name, lineno, "<type_key>.<field> ->
pub(crate)")
- self.visibility[_field_target(name, lhs, lineno)] = vis
- elif name == "enum":
- target, spec = _parse_enum(payload, lineno)
- self.enums[target] = spec
- elif name == "opaque":
- self.opaque.add(_type_target(name, payload, lineno))
- elif name == "upcast":
- lhs, rust_type = _split_arrow(name, payload, lineno, "<type_key>
-> <RustType>")
- self.upcasts.setdefault(_type_target(name, lhs, lineno),
[]).append(rust_type)
- elif name == "custom-new":
- self.custom_new.add(_type_target(name, payload, lineno))
- elif name == "no-new":
- self.no_new.add(_type_target(name, payload, lineno))
- elif name == "missing-docs":
- if payload.strip() != "allow": # the one level so far
- raise _invalid(name, lineno, "allow")
- self.allow_missing_docs = True
- else:
- raise ValueError(f"Unknown directive `{name}` at line {lineno}")
-
-
-def _invalid(name: str, lineno: int, expected: str) -> ValueError:
- return ValueError(f"Invalid `{name}` directive at line {lineno}. Expected
`{expected}`")
-
-
-def _type_target(name: str, text: str, lineno: int) -> str:
- """Validate a ``<type_key>`` reference."""
- target = text.strip()
- if not target or " " in target:
- raise _invalid(name, lineno, "<type_key>")
- return target
-
-
-def _field_target(name: str, text: str, lineno: int) -> str:
- """Validate a ``<type_key>.<field>`` reference."""
- target = text.strip()
- if not target or " " in target or "." not in target.strip("."):
- raise _invalid(name, lineno, "<type_key>.<field>")
- return target
-
-
-def _split_arrow(name: str, payload: str, lineno: int, expected: str) ->
tuple[str, str]:
- """Split ``<target> -> <rust type>``; the caller validates the target."""
- lhs, arrow, rhs = payload.partition("->")
- if not arrow or not rhs.strip():
- raise _invalid(name, lineno, expected)
- return lhs, rhs.strip()
-
-
-def _parse_enum(payload: str, lineno: int) -> tuple[str, EnumSpec]:
- """Parse ``<type_key>.<field> -> Name(i32) [closed] { A=0, B=1 }``.
-
- The member list is optional unless ``closed``, which needs a member equal
to 0.
- """
- expected = "<type_key>.<field> -> Name(i32) [closed] { A=0, B=1 }"
- match = _ENUM_RE.match(payload.strip())
- if match is None:
- raise _invalid("enum", lineno, expected)
- members: list[tuple[str, int]] = []
- for item in (match.group("body") or "").split(","):
- if not item.strip():
- continue
- member = _MEMBER_RE.match(item.strip())
- if member is None:
- raise _invalid("enum", lineno, expected)
- members.append((member.group("name"), int(member.group("value"))))
- closed = match.group("closed") is not None
- if closed and not members:
- raise _invalid("enum", lineno, expected)
- if closed and all(value != 0 for _, value in members):
+ @staticmethod
+ def from_directives(directives: Directives) -> RustDirectives:
+ """Collect the settings of ``directives``; raise ``ValueError`` on a
malformed ``enum``."""
+ settings, out = directives.settings, RustDirectives()
+ specs = {name: _enum_spec(name, keys) for name, keys in
settings["enum"].items()}
+ for target, keys in settings["field"].items():
+ rust_type = str(keys.get("type", ""))
+ if rust_type in specs:
+ out.enums[target] = specs[rust_type]
+ elif rust_type:
+ out.field_types[target] = rust_type
+ if "nullable" in keys:
+ out.nullable.add(target)
+ if "visibility" in keys:
+ out.visibility[target] = str(keys["visibility"])
+ for target, keys in settings["object"].items():
+ for name, chosen in (
+ ("opaque", out.opaque),
+ ("custom_new", out.custom_new),
+ ("no_new", out.no_new),
+ ):
+ if name in keys:
+ chosen.add(target)
+ upcasts = keys.get("upcast", [])
+ if isinstance(upcasts, list) and upcasts:
+ out.upcasts[target] = upcasts
+ out.allow_missing_docs = "missing_docs" in settings["file"].get("mod",
{})
+ return out
+
+
+def _enum_spec(name: str, keys: dict[str, bool | str | list[str]]) -> EnumSpec:
+ """Build the ``enum`` ``name``: a ``repr``, members ``Name = <int>``, and
``closed``."""
+ members = tuple((k, v) for k, v in keys.items() if k not in ("repr",
"closed"))
+ if "repr" not in keys:
+ raise ValueError(f"`enum` directive `{name}` needs `repr =
<{'|'.join(_INT_TYPES)}>`")
+ try:
+ values = tuple((member, int(str(value))) for member, value in members)
+ except ValueError:
+ raise ValueError(f"`enum` directive `{name}`: a member value is not an
integer") from None
+ closed = "closed" in keys
+ if closed and all(value != 0 for _, value in values):
raise ValueError(
- f"Closed `enum` directive at line {lineno} has no member equal to
0, the value of "
+ f"Closed `enum` directive `{name}` has no member equal to 0, the
value of "
"a field an object zero-filled at creation never set"
)
- target = _field_target("enum", match.group("target"), lineno)
- spec = EnumSpec(match.group("name"), match.group("repr"), tuple(members),
closed)
- return target, spec
+ return EnumSpec(name, str(keys["repr"]), values, closed)
diff --git a/python/tvm_ffi/stub/rust_generator/generator.py
b/python/tvm_ffi/stub/rust_generator/generator.py
index c7d0a92c..290f15a1 100644
--- a/python/tvm_ffi/stub/rust_generator/generator.py
+++ b/python/tvm_ffi/stub/rust_generator/generator.py
@@ -33,11 +33,13 @@ from .codegen import (
generate_rust_import_section,
generate_rust_object,
)
+from .directives import RUST_DIRECTIVE_KEYS, RustDirectives
from .utils import RustImports, RustUse, check_tokens, is_crate_type_key,
rust_module_path
if TYPE_CHECKING:
from pathlib import Path
+ from ..directives import Directives
from ..file_utils import CodeBlock
from ..utils import FuncInfo, InitConfig, ObjectInfo, Options
@@ -67,7 +69,7 @@ class RustGenerator:
name = "rust"
syntax = C.RUST_SYNTAX
source_exts = frozenset({".rs"})
- directive_kinds = C_RUST.RUST_DIRECTIVE_KINDS
+ directive_keys = RUST_DIRECTIVE_KEYS
def default_ty_map(self) -> dict[str, str]:
"""Return the default FFI-origin -> Rust-type name map."""
@@ -75,16 +77,12 @@ class RustGenerator:
# --- per-file collection
--------------------------------------------------
- def new_imports(self) -> RustImports:
- """Create an empty per-file collector (``use`` items and
directives)."""
- return RustImports()
-
- def add_directive(self, imports: RustImports, name: str, payload: str,
lineno: int) -> None:
- """Record ``import-object`` as a ``use`` (its path is the first ``;``
field); parse the rest."""
- if name == "import-object":
- imports.record(payload.split(";", 1)[0].strip())
- else:
- imports.directives.add(name, payload, lineno)
+ def new_imports(self, directives: Directives) -> RustImports:
+ """Create the per-file collector: the ``import`` paths as ``use``
items, and the settings."""
+ imports =
RustImports(directives=RustDirectives.from_directives(directives))
+ for path in directives.settings["import"]:
+ imports.record(path)
+ return imports
def canonical_type_name(self, type_key: str) -> str:
"""Return the Rust path for a defined type key (matches
:attr:`RustUse.path`)."""
diff --git a/python/tvm_ffi/stub/rust_generator/utils.py
b/python/tvm_ffi/stub/rust_generator/utils.py
index 7a5e0667..ec36bf90 100644
--- a/python/tvm_ffi/stub/rust_generator/utils.py
+++ b/python/tvm_ffi/stub/rust_generator/utils.py
@@ -25,7 +25,7 @@ from pathlib import Path
from typing import TYPE_CHECKING, Callable
from . import consts as C
-from .directives import Directives
+from .directives import RustDirectives
if TYPE_CHECKING:
from collections.abc import Iterable
@@ -105,7 +105,7 @@ class RustImports:
"""The per-file collector: the ``use`` items and the Rust directives of
one file."""
items: list[RustUse] = dataclasses.field(default_factory=list)
- directives: Directives = dataclasses.field(default_factory=Directives)
+ directives: RustDirectives =
dataclasses.field(default_factory=RustDirectives)
builtin_mirrors: dict[str, str] = dataclasses.field(default_factory=dict)
"""Builtin ancestors mirrored in this file, root first: type key -> the
base each embeds."""
reserved: set[str] = dataclasses.field(default_factory=set)
diff --git a/python/tvm_ffi/testing/testing.py
b/python/tvm_ffi/testing/testing.py
index b4797d67..eb847a04 100644
--- a/python/tvm_ffi/testing/testing.py
+++ b/python/tvm_ffi/testing/testing.py
@@ -161,7 +161,7 @@ class TestNonCopyable(Object):
@c_class("testing.SchemaAllTypes")
class _SchemaAllTypes:
- # tvm-ffi-stubgen(ty-map): testing.SchemaAllTypes ->
testing._SchemaAllTypes
+ # tvm-ffi-stubgen(object): testing.SchemaAllTypes(ty_map =
testing._SchemaAllTypes)
# tvm-ffi-stubgen(begin): object/testing.SchemaAllTypes
# fmt: off
v_bool: bool
@@ -317,7 +317,7 @@ class TestEqWithoutHash(Object):
@c_class("testing.TestCxxClassBase")
class _TestCxxClassBase(Object):
- # tvm-ffi-stubgen(ty-map): testing.TestCxxClassBase ->
testing._TestCxxClassBase
+ # tvm-ffi-stubgen(object): testing.TestCxxClassBase(ty_map =
testing._TestCxxClassBase)
# tvm-ffi-stubgen(begin): object/testing.TestCxxClassBase
# fmt: off
v_i64: int
@@ -338,7 +338,7 @@ class _TestCxxClassBase(Object):
@c_class("testing.TestCxxClassDerived", eq=True, order=True, unsafe_hash=True)
class _TestCxxClassDerived(_TestCxxClassBase):
- # tvm-ffi-stubgen(ty-map): testing.TestCxxClassDerived ->
testing._TestCxxClassDerived
+ # tvm-ffi-stubgen(object): testing.TestCxxClassDerived(ty_map =
testing._TestCxxClassDerived)
# tvm-ffi-stubgen(begin): object/testing.TestCxxClassDerived
# fmt: off
v_f64: float
@@ -352,7 +352,7 @@ class _TestCxxClassDerived(_TestCxxClassBase):
@c_class("testing.TestCxxClassDerivedDerived")
class _TestCxxClassDerivedDerived(_TestCxxClassDerived):
- # tvm-ffi-stubgen(ty-map): testing.TestCxxClassDerivedDerived ->
testing._TestCxxClassDerivedDerived
+ # tvm-ffi-stubgen(object): testing.TestCxxClassDerivedDerived(ty_map =
testing._TestCxxClassDerivedDerived)
# tvm-ffi-stubgen(begin): object/testing.TestCxxClassDerivedDerived
# fmt: off
v_str: str
@@ -396,7 +396,7 @@ class _TestCxxEnumHolder(Object):
@c_class("testing.TestCxxInitSubset")
class _TestCxxInitSubset(Object):
- # tvm-ffi-stubgen(ty-map): testing.TestCxxInitSubset ->
testing._TestCxxInitSubset
+ # tvm-ffi-stubgen(object): testing.TestCxxInitSubset(ty_map =
testing._TestCxxInitSubset)
# tvm-ffi-stubgen(begin): object/testing.TestCxxInitSubset
# fmt: off
required_field: int
@@ -411,7 +411,7 @@ class _TestCxxInitSubset(Object):
@c_class("testing.TestCxxKwOnly")
class _TestCxxKwOnly(Object):
- # tvm-ffi-stubgen(ty-map): testing.TestCxxKwOnly -> testing._TestCxxKwOnly
+ # tvm-ffi-stubgen(object): testing.TestCxxKwOnly(ty_map =
testing._TestCxxKwOnly)
# tvm-ffi-stubgen(begin): object/testing.TestCxxKwOnly
# fmt: off
x: int
@@ -431,7 +431,7 @@ class _TestCxxAutoInit(Object):
__test__ = False
- # tvm-ffi-stubgen(ty-map): testing.TestCxxAutoInit ->
testing._TestCxxAutoInit
+ # tvm-ffi-stubgen(object): testing.TestCxxAutoInit(ty_map =
testing._TestCxxAutoInit)
# tvm-ffi-stubgen(begin): object/testing.TestCxxAutoInit
# fmt: off
a: int
@@ -451,7 +451,7 @@ class _TestCxxAutoInitSimple(Object):
__test__ = False
- # tvm-ffi-stubgen(ty-map): testing.TestCxxAutoInitSimple ->
testing._TestCxxAutoInitSimple
+ # tvm-ffi-stubgen(object): testing.TestCxxAutoInitSimple(ty_map =
testing._TestCxxAutoInitSimple)
# tvm-ffi-stubgen(begin): object/testing.TestCxxAutoInitSimple
# fmt: off
x: int
@@ -469,7 +469,7 @@ class _TestCxxAutoInitAllInitOff(Object):
__test__ = False
- # tvm-ffi-stubgen(ty-map): testing.TestCxxAutoInitAllInitOff ->
testing._TestCxxAutoInitAllInitOff
+ # tvm-ffi-stubgen(object): testing.TestCxxAutoInitAllInitOff(ty_map =
testing._TestCxxAutoInitAllInitOff)
# tvm-ffi-stubgen(begin): object/testing.TestCxxAutoInitAllInitOff
# fmt: off
x: int
@@ -488,7 +488,7 @@ class _TestCxxAutoInitKwOnlyDefaults(Object):
__test__ = False
- # tvm-ffi-stubgen(ty-map): testing.TestCxxAutoInitKwOnlyDefaults ->
testing._TestCxxAutoInitKwOnlyDefaults
+ # tvm-ffi-stubgen(object): testing.TestCxxAutoInitKwOnlyDefaults(ty_map =
testing._TestCxxAutoInitKwOnlyDefaults)
# tvm-ffi-stubgen(begin): object/testing.TestCxxAutoInitKwOnlyDefaults
# fmt: off
p_required: int
@@ -509,7 +509,7 @@ class _TestCxxNoAutoInit(Object):
__test__ = False
- # tvm-ffi-stubgen(ty-map): testing.TestCxxNoAutoInit ->
testing._TestCxxNoAutoInit
+ # tvm-ffi-stubgen(object): testing.TestCxxNoAutoInit(ty_map =
testing._TestCxxNoAutoInit)
# tvm-ffi-stubgen(begin): object/testing.TestCxxNoAutoInit
# fmt: off
x: int
@@ -524,7 +524,7 @@ class _TestCxxAutoInitParent(Object):
__test__ = False
- # tvm-ffi-stubgen(ty-map): testing.TestCxxAutoInitParent ->
testing._TestCxxAutoInitParent
+ # tvm-ffi-stubgen(object): testing.TestCxxAutoInitParent(ty_map =
testing._TestCxxAutoInitParent)
# tvm-ffi-stubgen(begin): object/testing.TestCxxAutoInitParent
# fmt: off
parent_required: int
@@ -542,7 +542,7 @@ class _TestCxxAutoInitChild(_TestCxxAutoInitParent):
__test__ = False
- # tvm-ffi-stubgen(ty-map): testing.TestCxxAutoInitChild ->
testing._TestCxxAutoInitChild
+ # tvm-ffi-stubgen(object): testing.TestCxxAutoInitChild(ty_map =
testing._TestCxxAutoInitChild)
# tvm-ffi-stubgen(begin): object/testing.TestCxxAutoInitChild
# fmt: off
child_required: int
diff --git a/tests/python/test_stubgen.py b/tests/python/test_stubgen.py
index 0de76e91..aba0aeee 100644
--- a/tests/python/test_stubgen.py
+++ b/tests/python/test_stubgen.py
@@ -27,6 +27,7 @@ from tvm_ffi.core import MISSING, TypeSchema,
_lookup_or_register_type_info_from
from tvm_ffi.dataclasses import py_class
from tvm_ffi.stub import consts as C
from tvm_ffi.stub.cli import _stage_1, _stage_2, _stage_3
+from tvm_ffi.stub.directives import Directives, Key
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
@@ -45,7 +46,7 @@ from tvm_ffi.stub.python_generator.codegen import (
render_object_init,
render_object_methods,
)
-from tvm_ffi.stub.python_generator.utils import ImportItem
+from tvm_ffi.stub.python_generator.utils import ImportItem, PythonImports
from tvm_ffi.stub.utils import (
FuncInfo,
InitConfig,
@@ -101,10 +102,10 @@ 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"
+ line = f"{C.PYTHON_SYNTAX.directive('object')} custom(ty_map = mapped)"
block = CodeBlock.from_begin_line(5, line, C.PYTHON_SYNTAX)
assert block.kind == "directive"
- assert block.param == ("ty-map", "custom -> mapped")
+ assert block.param == ("object", "custom(ty_map = mapped)")
assert block.lineno_start == 5
assert block.lineno_end == 5
@@ -129,7 +130,7 @@ def test_fileinfo_from_file_parses_blocks(tmp_path: Path)
-> None:
f"{C.PYTHON_SYNTAX.begin} global/demo.func",
"in_stub = True",
C.PYTHON_SYNTAX.end,
- f"{C.PYTHON_SYNTAX.ty_map} x -> y",
+ f"{C.PYTHON_SYNTAX.directive('object')} x(ty_map = y)",
]
)
path = tmp_path / "demo.py"
@@ -154,9 +155,9 @@ def test_fileinfo_from_file_parses_blocks(tmp_path: Path)
-> None:
]
assert ty_map.kind == "directive"
- assert ty_map.param == ("ty-map", "x -> y")
+ assert ty_map.param == ("object", "x(ty_map = y)")
assert ty_map.lineno_start == ty_map.lineno_end == 5
- assert ty_map.lines == [f"{C.PYTHON_SYNTAX.ty_map} x -> y"]
+ assert ty_map.lines == [f"{C.PYTHON_SYNTAX.directive('object')} x(ty_map =
y)"]
def test_fileinfo_from_file_error_paths(tmp_path: Path) -> None:
@@ -835,11 +836,12 @@ def
test_stage_3_adds_LIB_when_load_lib_imported(tmp_path: Path) -> None:
)
import_obj_block = CodeBlock(
kind="directive",
- param=("import-object",
"tvm_ffi.libinfo.load_lib_module;False;_FFI_LOAD_LIB"),
+ param=("import", "tvm_ffi.libinfo.load_lib_module(alias =
_FFI_LOAD_LIB)"),
lineno_start=1,
lineno_end=1,
lines=[
- f"{C.PYTHON_SYNTAX.import_object}
tvm_ffi.libinfo.load_lib_module;False;_FFI_LOAD_LIB"
+ f"{C.PYTHON_SYNTAX.directive('import')} "
+ "tvm_ffi.libinfo.load_lib_module(alias = _FFI_LOAD_LIB)"
],
)
all_block = CodeBlock(
@@ -945,11 +947,12 @@ def test_generate_ffi_api_with_objects_imports_parents()
-> None:
is_root=False,
syntax=C.PYTHON_SYNTAX,
)
- assert C.PYTHON_SYNTAX.import_object in code # register_object prompt
+ assert C.PYTHON_SYNTAX.directive("import") in code # register_object
prompt
assert f"{C.PYTHON_SYNTAX.begin} object/{obj_info.type_key}" in code
assert parent_key is not None
parent_import_prompt = (
- f"{C.PYTHON_SYNTAX.import_object}
{parent_key};False;_{parent_key.replace('.', '_')}"
+ f"{C.PYTHON_SYNTAX.directive('import')} {parent_key}"
+ f"(alias = _{parent_key.replace('.', '_')})"
)
assert parent_import_prompt in code
@@ -1095,7 +1098,7 @@ def test_rust_marker_syntax_parses_rs_file(tmp_path:
Path) -> None:
f"{C.RUST_SYNTAX.begin} object/demo.Foo",
"pub struct Foo;",
C.RUST_SYNTAX.end,
- f"{C.RUST_SYNTAX.ty_map} a -> b",
+ f"{C.RUST_SYNTAX.directive('object')} a(ty_map = b)",
"",
]
),
@@ -1107,7 +1110,7 @@ def test_rust_marker_syntax_parses_rs_file(tmp_path:
Path) -> None:
assert info.syntax is C.RUST_SYNTAX
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[1].param == ("object", "a(ty_map = b)")
assert info.code_blocks[0].lines[1] == "pub struct Foo;"
# Python markers are plain text inside a Rust file.
@@ -1141,15 +1144,17 @@ def test_generator_registry_names() -> None:
def test_codeblock_from_begin_line_directive() -> None:
- """``<comment> tvm-ffi-stubgen(<name>): <payload>`` parses into a one-line
directive block."""
+ """``<comment> tvm-ffi-stubgen(<subject>): <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)
+ block = CodeBlock.from_begin_line(
+ 7, f"{syntax.directive('field')} a.B.c(type = D) ", syntax
+ )
assert block.kind == "directive"
- assert block.param == ("field", "a.B.c -> D")
+ assert block.param == ("field", "a.B.c(type = 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", "")
+ empty = CodeBlock.from_begin_line(1, C.PYTHON_SYNTAX.directive("object"),
C.PYTHON_SYNTAX)
+ assert empty.param == ("object", "")
# 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"):
@@ -1161,7 +1166,7 @@ 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",
+ " // tvm-ffi-stubgen(field): ir.Expr.span(nullable)",
f"{C.RUST_SYNTAX.begin} object/demo.Foo",
C.RUST_SYNTAX.end,
f"{C.RUST_SYNTAX.prefix}bogus)",
@@ -1174,28 +1179,28 @@ def
test_fileinfo_from_file_keeps_directive_lines(tmp_path: Path) -> None:
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].param == ("field", "ir.Expr.span(nullable)")
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."""
+ """Pipeline keys stay in the pipeline; declared subjects and keys reach
the generator."""
path = tmp_path / "demo.py"
ty_map_directive = CodeBlock(
kind="directive",
- param=("ty-map", "demo.Foo -> mapped.Foo"),
+ param=("object", "demo.Foo(ty_map = mapped.Foo)"),
lineno_start=1,
lineno_end=1,
- lines=[f"{C.PYTHON_SYNTAX.ty_map} demo.Foo -> mapped.Foo"],
+ lines=[f"{C.PYTHON_SYNTAX.directive('object')} demo.Foo(ty_map =
mapped.Foo)"],
)
directive = CodeBlock(
kind="directive",
- param=("field", "demo.Foo.x -> Bar"),
+ param=("field", "demo.Foo.x(type = Bar)"),
lineno_start=2,
lineno_end=2,
- lines=[f"{C.PYTHON_SYNTAX.directive('field')} demo.Foo.x -> Bar"],
+ lines=[f"{C.PYTHON_SYNTAX.directive('field')} demo.Foo.x(type = Bar)"],
)
all_block = CodeBlock(
kind="__all__",
@@ -1214,26 +1219,28 @@ def
test_stage_3_routes_directives_by_generator_declaration(tmp_path: Path) -> N
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
+ seen: list[Directives] = []
+
+ keys = {**PythonGenerator.directive_keys, "field": {"type": Key(True)}}
+
+ class FieldAware(PythonGenerator):
+ directive_keys = keys
+
+ def new_imports(self, directives: Directives) -> PythonImports:
+ seen.append(directives)
+ return super().new_imports(directives)
+
+ # `ty_map` is a pipeline key: `_stage_1` reads it.
+ assert "ty_map" in C.PIPELINE_DIRECTIVE_KEYS["object"]
ty_map: dict[str, str] = {}
- _stage_1(_file_info(), ty_map)
+ _stage_1(_file_info(), ty_map, FieldAware())
assert ty_map == {"demo.Foo": "mapped.Foo"}
- # The Python generator does not declare `field`, so the name is unknown to
it.
+ # The Python generator does not declare `field`, so the subject is unknown
to it.
python = get_generator("python")
- assert python.directive_kinds == frozenset({"import-object"})
+ assert set(python.directive_keys) == {"import"}
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()))
+ assert seen[0].settings["field"] == {"demo.Foo.x": {"type": "Bar"}}
diff --git a/tests/python/test_stubgen_rust.py
b/tests/python/test_stubgen_rust.py
index b257de53..ae9cd487 100644
--- a/tests/python/test_stubgen_rust.py
+++ b/tests/python/test_stubgen_rust.py
@@ -29,6 +29,7 @@ import tvm_ffi.testing # loads the `testing.*` fixture types
from tvm_ffi.core import TypeSchema
from tvm_ffi.stub import consts as C
from tvm_ffi.stub.cli import _stage_1, _stage_3
+from tvm_ffi.stub.directives import Directives
from tvm_ffi.stub.file_utils import CodeBlock, FileInfo
from tvm_ffi.stub.generator import get_generator
from tvm_ffi.stub.lib_state import object_info_from_type_key
@@ -40,7 +41,7 @@ from tvm_ffi.stub.rust_generator.codegen import (
generate_rust_import_section,
generate_rust_object,
)
-from tvm_ffi.stub.rust_generator.directives import Directives, EnumSpec
+from tvm_ffi.stub.rust_generator.directives import RUST_DIRECTIVE_KEYS,
EnumSpec
from tvm_ffi.stub.rust_generator.utils import (
RustImports,
RustUse,
@@ -298,16 +299,28 @@ def test_check_tokens_ignores_spelling() -> None:
# ---------------------------------------------------------------------------
+def _imports(*lines: str) -> RustImports:
+ """Seed a collector from directive lines ``<subject>: <payload>``,
numbered from 1."""
+ directives = Directives(RUST_DIRECTIVE_KEYS)
+ for lineno, line in enumerate(lines, 1):
+ subject, _, payload = line.partition(": ")
+ directives.add(subject, payload, lineno)
+ return RUST.new_imports(directives)
+
+
def test_directives_parse() -> None:
- directives = Directives()
- directives.add("field", " tirx.Add.a -> PrimExpr ", 1)
- directives.add("nullable", "ir.Expr.span", 2)
- directives.add("enum", "tirx.For.kind -> ForKind(i32) { Serial=0, Parallel
= 1 }", 3)
- directives.add("enum", "tirx.For.mode -> Mode(u8)", 4)
- directives.add("opaque", " ir.SourceName ", 5)
- directives.add("upcast", "tirx.Add -> PrimExpr", 6)
- directives.add("upcast", "tirx.Add -> crate::typed::TypedExpr", 7)
- directives.add("custom-new", " tirx.Add ", 8)
+ directives = _imports(
+ "field: tirx.Add.a( type = PrimExpr )",
+ "field: ir.Expr.span(nullable)",
+ "enum: ForKind(repr = i32, Serial = 0, Parallel = 1)",
+ "field: tirx.For.kind(type = ForKind)",
+ "enum: Mode(repr = u8)",
+ "field: tirx.For.mode(type = Mode)",
+ "object: ir.SourceName(opaque) ",
+ "object: tirx.Add(upcast = PrimExpr)",
+ "object: tirx.Add(upcast = crate::typed::TypedExpr, custom_new)",
+ "object: tirx.Add",
+ ).directives
assert directives.field_types == {"tirx.Add.a": "PrimExpr"}
assert directives.nullable == {"ir.Expr.span"}
assert directives.enums == {
@@ -320,44 +333,35 @@ def test_directives_parse() -> None:
@pytest.mark.parametrize(
- ("name", "payload", "expected"),
+ ("line", "expected"),
[
- ("field", "tirx.Add.a", "-> <RustType>"),
- ("field", "tirx.Add.a ->", "-> <RustType>"),
- ("field", "Add -> PrimExpr", "<type_key>.<field>"),
- ("nullable", "ir.Expr.span extra", "<type_key>.<field>"),
- ("enum", "tirx.For.kind -> ForKind", "Name(i32)"),
- ("enum", "tirx.For.kind -> ForKind(i128)", "Name(i32)"),
- ("enum", "tirx.For.kind -> ForKind(i32) { Serial }", "Name(i32)"),
- ("opaque", "ir.SourceName ir.Source", "<type_key>"),
- ("upcast", "tirx.Add", "-> <RustType>"),
- ("upcast", "tirx.Add PrimExpr -> PrimExpr", "<type_key>"),
- ("custom-new", "", "<type_key>"),
- ("typed-view", "tirx.Add -> PrimExpr", "Unknown directive"),
+ ("field: tirx.Add.a(type)", "`type` takes a value"),
+ ("field: tirx.Add.a(type =)", "`type` takes a value"),
+ ("field: tirx.Add.a(type = PrimType)", "`type` conflicts with line 6"),
+ ("field: Add(type = PrimExpr)", "<type_key>.<field>"),
+ ("field: ir.Expr.span extra", "bad target"),
+ ("field: ir.Expr.span(nullable", "expected `<target>("),
+ ("enum: ForKind(repr = i128)", "`repr` is one of"),
+ ("enum: ForKind(repr = i32, Serial)", "`Serial` takes a value"),
+ ("object: ir.SourceName ir.Source", "bad target"),
+ ("object: tirx.Add(upcast)", "`upcast` takes a value"),
+ ("object: tirx.Add(custom_new = yes)", "`custom_new` is a flag"),
+ ("object: tirx.Add(typed_view = PrimExpr)", "unknown key
`typed_view`"),
+ ("typed-view: tirx.Add", "Unknown directive"),
],
)
-def test_directives_reject_malformed(name: str, payload: str, expected: str)
-> None:
+def test_directives_reject_malformed(line: str, expected: str) -> None:
+ directives = Directives(RUST_DIRECTIVE_KEYS)
+ directives.add("field", "tirx.Add.a(type = PrimExpr)", 6)
+ subject, _, payload = line.partition(": ")
with pytest.raises(ValueError, match=re.escape(expected)) as exc:
- Directives().add(name, payload, 7)
+ directives.add(subject, payload, 7)
assert "at line 7" in str(exc.value)
def test_generator_declares_its_directives_and_records_imports() -> None:
- assert RUST.directive_kinds == {
- "import-object",
- "field",
- "nullable",
- "visibility",
- "enum",
- "opaque",
- "upcast",
- "custom-new",
- "no-new",
- "missing-docs",
- }
- imports = RUST.new_imports()
- RUST.add_directive(imports, "import-object",
"tvm_ffi.libinfo.Foo;False;_Foo", 1)
- RUST.add_directive(imports, "nullable", "demo.Node.span", 2)
+ assert set(RUST.directive_keys) == {"object", "field", "enum", "file",
"import"}
+ imports = _imports("import: tvm_ffi::libinfo::Foo", "field:
demo.Node.span(nullable)")
assert imports.items == [RustUse("tvm_ffi::libinfo::Foo")]
assert imports.directives.nullable == {"demo.Node.span"}
@@ -640,15 +644,12 @@ def test_render_iter_var_golden() -> None:
parent="ir.PrimExprConvertible",
is_final=True,
)
- imports = RUST.new_imports()
- RUST.add_directive(imports, "nullable", "tirx.IterVar.dom", 1)
- RUST.add_directive(imports, "field", "tirx.IterVar.var -> PrimVar", 2)
- RUST.add_directive(imports, "nullable", "tirx.IterVar.span", 3)
- RUST.add_directive(
- imports,
- "enum",
- "tirx.IterVar.iter_type -> IterVarType(i32) { kDataPar=0,
kThreadIndex=1 }",
- 4,
+ imports = _imports(
+ "field: tirx.IterVar.dom(nullable)",
+ "field: tirx.IterVar.var(type = PrimVar)",
+ "field: tirx.IterVar.span(nullable)",
+ "enum: IterVarType(repr = i32, kDataPar = 0, kThreadIndex = 1)",
+ "field: tirx.IterVar.iter_type(type = IterVarType)",
)
text, imports = _render(info, imports)
assert text == ITER_VAR_EXPECTED
@@ -668,9 +669,10 @@ def
test_field_directive_with_path_records_use_and_nullable_does_not_double_wrap
_field("dom", TypeSchema("Optional", (TypeSchema("int"),))),
),
)
- imports = RUST.new_imports()
- RUST.add_directive(imports, "field", "demo.Node.buffer ->
crate::typed::BufferVar", 1)
- RUST.add_directive(imports, "nullable", "demo.Node.dom", 2)
+ imports = _imports(
+ "field: demo.Node.buffer(type = crate::typed::BufferVar)",
+ "field: demo.Node.dom(nullable)",
+ )
text, imports = _render(info, imports)
assert "pub fn buffer(&self) -> Result<BufferVar> {" in text
assert "pub fn dom(&self) -> Result<Option<i64>> {" in text
@@ -696,12 +698,10 @@ def _span() -> ObjectInfo:
)
-def _span_imports() -> RustImports:
+def _span_imports(*extra: str) -> RustImports:
"""Type the narrow integer fields by directive: the schema erases
signedness."""
- imports = RUST.new_imports()
- for lineno, name in enumerate(("line", "column", "end_line",
"end_column"), 1):
- RUST.add_directive(imports, "field", f"ir.Span.{name} -> i32", lineno)
- return imports
+ names = ("line", "column", "end_line", "end_column")
+ return _imports(*(f"field: ir.Span.{name}(type = i32)" for name in names),
*extra)
def _expr() -> ObjectInfo:
@@ -816,8 +816,7 @@ tvm_ffi::impl_object_ref_upcast!(unsafe Add => PrimExpr);"""
def test_render_complete_expr_golden() -> None:
"""`ir.Expr` as tvm-rust-ext hand-writes it: `span: Option<Span>`, `ty:
Type`, no getters."""
- imports = RUST.new_imports()
- RUST.add_directive(imports, "nullable", "ir.Expr.span", 1)
+ imports = _imports("field: ir.Expr.span(nullable)")
text, imports = _render(_expr(), imports)
assert text == EXPR_EXPECTED
assert _uses(imports) == {"tvm_ffi::Object", "tvm_ffi::IntoArg",
"tvm_ffi::ObjectArc"}
@@ -831,12 +830,13 @@ def test_render_complete_add_golden() -> None:
as it is); `upcast` adds the `PrimExpr` view.
"""
_register(_expr())
- imports = RUST.new_imports()
- RUST.add_directive(imports, "nullable", "ir.Expr.span", 1)
- RUST.add_directive(imports, "field", "tirx.Add.a -> PrimExpr", 2)
- RUST.add_directive(imports, "field", "tirx.Add.b -> PrimExpr", 3)
- RUST.add_directive(imports, "field", "tirx.Add.ty -> PrimType", 4)
- RUST.add_directive(imports, "upcast", "tirx.Add -> PrimExpr", 5)
+ imports = _imports(
+ "field: ir.Expr.span(nullable)",
+ "field: tirx.Add.a(type = PrimExpr)",
+ "field: tirx.Add.b(type = PrimExpr)",
+ "field: tirx.Add.ty(type = PrimType)",
+ "object: tirx.Add(upcast = PrimExpr)",
+ )
text, imports = _render(_add(), imports)
assert text == ADD_EXPECTED
assert _uses(imports) == {
@@ -895,23 +895,21 @@ def test_reflected_base_and_data_keep_their_names() ->
None:
total_size=48,
is_final=True,
)
- imports = RUST.new_imports()
- RUST.add_directive(imports, "visibility", "demo.Ramp.cache_ ->
pub(crate)", 1)
+ imports = _imports("field: demo.Ramp.cache_(visibility = pub(crate))")
text, _ = _render(info, imports)
assert (
" base_: Object,\n pub base: i64,\n pub data: i64,\n
pub(crate) cache_: i64,\n"
) in text
assert " pub fn new(base: i64, data: i64, cache_: i64) -> Self {" in
text
assert " Self { base_, base, data, cache_ }" in text
- imports = RUST.new_imports()
- RUST.add_directive(imports, "visibility", "demo.Node.cache_ ->
pub(crate)", 1)
+ imports = _imports("field: demo.Node.cache_(visibility = pub(crate))")
text, _ = _render(_info("demo.Node", (_field("base", "int"),
_field("cache_", "int"))), imports)
assert (
'pub fn base(&self) -> Result<i64> {\n
cached_field_getter!(NodeObj, "base")' in text
)
assert " pub(crate) fn cache_(&self) -> Result<i64> {" in text
- with pytest.raises(ValueError, match=r"`visibility` directive at line 2"):
- Directives().add("visibility", "demo.Node.cache_ -> pub(super)", 2)
+ with pytest.raises(ValueError, match=r"line 2: `visibility` is one of
pub\(crate\)"):
+ Directives(RUST_DIRECTIVE_KEYS).add("field",
"demo.Node.cache_(visibility = pub(super))", 2)
def test_render_complete_enum_field() -> None:
@@ -922,8 +920,7 @@ def test_render_complete_enum_field() -> None:
total_size=40,
is_final=True,
)
- imports = RUST.new_imports()
- RUST.add_directive(imports, "enum", "demo.Pair.kind -> Kind(i32) { A=0,
B=1 }", 1)
+ imports = _imports("enum: Kind(repr = i32, A = 0, B = 1)", "field:
demo.Pair.kind(type = Kind)")
text, imports = _render(info, imports)
assert " pub kind: Kind,\n" in text
assert " pub fn new(a: i64, kind: impl Into<Kind>) -> Self {" in text
@@ -932,16 +929,18 @@ def test_render_complete_enum_field() -> None:
def test_allow_missing_docs_and_closed_enum() -> None:
- """`missing-docs: allow` marks every generated type and `impl`; `closed`
emits a Rust enum."""
+ """`missing_docs = allow` marks every generated type and `impl`; `closed`
emits a Rust enum."""
info = _info(
"demo.Pair",
(_field("a", "int", 24, 8), _field("kind", "int", 32, 4)),
total_size=40,
is_final=True,
)
- imports = RUST.new_imports()
- RUST.add_directive(imports, "missing-docs", " allow ", 1)
- RUST.add_directive(imports, "enum", "demo.Pair.kind -> Kind(i32) closed {
A=0, B=-1 }", 2)
+ imports = _imports(
+ "file: mod(missing_docs = allow)",
+ "enum: Kind(repr = i32, closed, A = 0, B = -1)",
+ "field: demo.Pair.kind(type = Kind)",
+ )
text, _ = _render(info, imports)
assert text.startswith(
"#[allow(missing_docs)]\n"
@@ -962,15 +961,12 @@ def test_allow_missing_docs_and_closed_enum() -> None:
) in text
# The enum, the object struct, the wrapper and both allocator `impl`s.
assert text.count("#[allow(missing_docs)]\n") == 5
- with pytest.raises(ValueError, match="closed"):
- Directives().add("enum", "demo.Pair.kind -> Kind(i32) closed {}", 3)
- with pytest.raises(ValueError, match="no member equal to 0"):
- Directives().add("enum", "demo.Pair.kind -> Kind(i32) closed { A=1 }",
4)
- for payload in ("", "warn"):
- with pytest.raises(
- ValueError, match=r"`missing-docs` directive at line 4\. Expected
`allow`"
- ):
- Directives().add("missing-docs", payload, 4)
+ for members in ("", ", A = 1"):
+ with pytest.raises(ValueError, match="no member equal to 0"):
+ _imports(f"enum: Kind(repr = i32, closed{members})")
+ for payload, expected in (("", "takes a value"), (" = warn", "is one of
allow")):
+ with pytest.raises(ValueError, match=f"line 4: `missing_docs`
{expected}"):
+ Directives(RUST_DIRECTIVE_KEYS).add("file",
f"mod(missing_docs{payload})", 4)
def test_complete_optional_field_mirrors() -> None:
@@ -1046,7 +1042,7 @@ def test_complete_big_int_field_mirrors() -> None:
def test_ty_map_origin_override_survives_lossless_fold() -> None:
- """A `ty-map` for `str` still wins for a reflected `ffi.String`: the
lossless fold keeps no key."""
+ """A `ty_map` for `str` still wins for a reflected `ffi.String`: the
lossless fold keeps no key."""
name = TypeSchema.from_json_str('{"type":"ffi.String"}')
assert (name.origin, name.type_key) == ("str", None)
info = _info("demo.Named", (_field("name", name, 24, 16),), total_size=40,
is_final=True)
@@ -1149,10 +1145,9 @@ def test_unrenderable_field_keeps_the_type_opaque(field:
NamedTypeSchema) -> Non
def test_custom_new_renames_the_wrapper_allocator() -> None:
- """`custom-new`: `Add::new` stays hand-written, the allocator is
`from_fields`."""
+ """`custom_new`: `Add::new` stays hand-written, the allocator is
`from_fields`."""
_register(_expr())
- imports = RUST.new_imports()
- RUST.add_directive(imports, "custom-new", "tirx.Add", 1)
+ imports = _imports("object: tirx.Add(custom_new)")
text, _ = _render(_add(), imports)
assert (
"impl AddObj {\n pub(crate) fn new(span: Span, ty: Type, a: Expr,
b: Expr) -> Self {"
@@ -1203,11 +1198,13 @@ def
test_wrapper_allocator_takes_object_references_as_into_arg() -> None:
total_size=160,
is_final=True,
)
- imports = RUST.new_imports()
- RUST.add_directive(imports, "field", "demo.Node.buffer ->
crate::typed::BufferVar", 1)
- RUST.add_directive(imports, "nullable", "demo.Node.span", 2)
- RUST.add_directive(imports, "enum", "demo.Node.kind -> Kind(i64) { A=0 }",
3)
- RUST.add_directive(imports, "field", "demo.Node.count -> i32", 4)
+ imports = _imports(
+ "field: demo.Node.buffer(type = crate::typed::BufferVar)",
+ "field: demo.Node.span(nullable)",
+ "enum: Kind(repr = i64, A = 0)",
+ "field: demo.Node.kind(type = Kind)",
+ "field: demo.Node.count(type = i32)",
+ )
text, imports = _render(info, imports)
assert (
" pub fn new(\n"
@@ -1284,9 +1281,10 @@ def test_allocator_body_wraps_like_rustfmt() -> None:
def test_upcast_directive_adds_typed_views() -> None:
"""`upcast` targets follow the ancestor chain; a `::` path is imported.
Opaque: no allocator."""
info = _info("demo.Leaf", parent="demo.Base")
- imports = RUST.new_imports()
- RUST.add_directive(imports, "upcast", "demo.Leaf ->
crate::typed::LeafView", 1)
- RUST.add_directive(imports, "upcast", "demo.Leaf -> Other", 2)
+ imports = _imports(
+ "object: demo.Leaf(upcast = crate::typed::LeafView)",
+ "object: demo.Leaf(upcast = Other)",
+ )
text, imports = _render(info, imports)
assert text.endswith(
"tvm_ffi::impl_object_ref_upcast!(Leaf => Base);\n"
@@ -1297,8 +1295,7 @@ def test_upcast_directive_adds_typed_views() -> None:
def test_opaque_directive_vetoes_a_complete_type() -> None:
- imports = _span_imports()
- RUST.add_directive(imports, "opaque", "ir.Span", 5)
+ imports = _span_imports("object: ir.Span(opaque)")
text, _ = _render(_span(), imports)
assert "// Opaque: vetoed by directive although the layout is
reproducible." in text
assert "pub fn line(&self) -> Result<i32> {" in text
@@ -1335,19 +1332,19 @@ def test_opaque_parent_keeps_the_child_opaque() -> None:
@pytest.mark.parametrize(
- ("name", "payload", "message"),
+ ("lines", "message"),
[
- ("field", "demo.Pair.count -> i64", "maps a 4-byte field to `i64` (8
bytes)"),
- ("enum", "demo.Pair.count -> Kind(i64)", "maps a 4-byte field to `i64`
(8 bytes)"),
- ("nullable", "demo.Pair.count", "the field is 4 bytes, not a
pointer-sized"),
+ (("field: demo.Pair.count(type = i64)",), "maps a 4-byte field to
`i64` (8 bytes)"),
+ (
+ ("enum: Kind(repr = i64)", "field: demo.Pair.count(type = Kind)"),
+ "maps a 4-byte field to `i64` (8 bytes)",
+ ),
+ (("field: demo.Pair.count(nullable)",), "the field is 4 bytes, not a
pointer-sized"),
],
)
-def test_directive_disagreeing_with_bytes_is_an_error(
- name: str, payload: str, message: str
-) -> None:
+def test_directive_disagreeing_with_bytes_is_an_error(lines: tuple[str, ...],
message: str) -> None:
info = _info("demo.Pair", (_field("count", "int", 24, 4),), total_size=32)
- imports = RUST.new_imports()
- RUST.add_directive(imports, name, payload, 1)
+ imports = _imports(*lines)
with pytest.raises(ValueError, match=re.escape(message)):
_render(info, imports)
@@ -1407,8 +1404,7 @@ def test_finalize_module_tree(tmp_path: Path) -> None:
def _cxx_class_imports() -> RustImports:
"""Type the narrow integer ``v_i32`` by directive."""
- imports = RUST.new_imports()
- RUST.add_directive(imports, "field", "testing.TestCxxClassBase.v_i32 ->
i32", 1)
+ imports = _imports("field: testing.TestCxxClassBase.v_i32(type = i32)")
return imports
@@ -1483,7 +1479,8 @@ def
test_stage_3_applies_directives_to_a_registered_type(tmp_path: Path) -> None
[
f"{C.RUST_SYNTAX.begin} import-section",
C.RUST_SYNTAX.end,
- f"{C.RUST_SYNTAX.directive('enum')}
testing.TestCxxClassBase.v_i32 -> Kind(i32) {{ A=0 }}",
+ f"{C.RUST_SYNTAX.directive('enum')} Kind(repr = i32, A = 0)",
+ f"{C.RUST_SYNTAX.directive('field')}
testing.TestCxxClassBase.v_i32(type = Kind)",
f"{C.RUST_SYNTAX.begin} object/testing.TestCxxClassBase",
C.RUST_SYNTAX.end,
"",
@@ -1799,9 +1796,7 @@ def test_allocator_hands_a_ty_map_parent_its_own_types()
-> None:
is_final=True,
)
ty_map = {**RUST.default_ty_map(), "demo.Base": "crate::hand::Base"}
- imports = RUST.new_imports()
- RUST.add_directive(imports, "field", "demo.Base.kind -> i32", 1)
- RUST.add_directive(imports, "field", "demo.Child.kind -> Kind", 2)
+ imports = _imports("field: demo.Base.kind(type = i32)", "field:
demo.Child.kind(type = Kind)")
block = _object_block("demo.Child")
generate_rust_object(block, ty_map, imports, Options(), info,
_declared("demo.Child"))
text = "\n".join(block.lines[1:-1])
@@ -1833,14 +1828,17 @@ def
test_stage_3_checks_dependencies_across_the_run(tmp_path: Path) -> None:
assert info is not None
with pytest.raises(ValueError, match=r"depends on
`testing\.TestCxxClassBase`"):
_stage_3(info, Options(dry_run=True), RUST.default_ty_map(), {}, RUST,
declared)
- # A `ty-map` names a hand-written parent; its `<Name>Obj` becomes the base
slot.
- mapped = f"{C.RUST_SYNTAX.ty_map} testing.TestCxxClassBase ->
crate::hand::TestCxxClassBase"
- narrow = f"{C.RUST_SYNTAX.directive('field')}
testing.TestCxxClassBase.v_i32 -> i32"
+ # A `ty_map` names a hand-written parent; its `<Name>Obj` becomes the base
slot.
+ mapped = (
+ f"{C.RUST_SYNTAX.directive('object')} "
+ "testing.TestCxxClassBase(ty_map = crate::hand::TestCxxClassBase)"
+ )
+ narrow = f"{C.RUST_SYNTAX.directive('field')}
testing.TestCxxClassBase.v_i32(type = i32)"
src.write_text("\n".join([mapped, narrow, *blocks]), encoding="utf-8")
info = FileInfo.from_file(src)
assert info is not None
ty_map = RUST.default_ty_map()
- _stage_1(info, ty_map)
+ _stage_1(info, ty_map, RUST)
_stage_3(info, Options(dry_run=True), ty_map, {}, RUST, declared)
text = "\n".join(line for block in info.code_blocks for line in
block.lines)
assert " base_: TestCxxClassBaseObj," in text
@@ -1856,7 +1854,7 @@ def
test_stage_3_checks_dependencies_across_the_run(tmp_path: Path) -> None:
def test_ffi_field_type_off_the_default_ty_map_reads_as_any() -> None:
- """Of the `ffi.*` keys, only those on the default `ty-map` name a type the
crate binds."""
+ """Of the `ffi.*` keys, only those on the default `ty_map` name a type the
crate binds."""
text, imports = _render(_info("demo.Doc", (_field("path",
"ffi.reflection.AccessPath"),)))
assert " pub fn path(&self) -> Result<AccessPath> {" in text
assert "tvm_ffi::reflection::AccessPath" in _uses(imports)
@@ -1935,8 +1933,8 @@ def test_prefix_directive_rolls_out_a_namespace(
"\n".join(
[
"//! Hand-written skeleton.",
- f"{C.RUST_SYNTAX.directive('prefix')} testing",
- f"{C.RUST_SYNTAX.directive('skip')}
testing.TestCxxClassDerivedDerived",
+ f"{C.RUST_SYNTAX.directive('file')} mod(prefix = testing)",
+ f"{C.RUST_SYNTAX.directive('object')}
testing.TestCxxClassDerivedDerived(skip)",
"",
f"{C.RUST_SYNTAX.begin} object/testing.TestCxxClassBase",
C.RUST_SYNTAX.end,
@@ -1973,7 +1971,9 @@ def test_prefix_declared_twice_is_a_failure(
) -> None:
root = _crate(tmp_path)
for name in ("a.rs", "b.rs"):
- (root / name).write_text(f"{C.RUST_SYNTAX.directive('prefix')}
testing\n", encoding="utf-8")
+ (root / name).write_text(
+ f"{C.RUST_SYNTAX.directive('file')} mod(prefix = testing)\n",
encoding="utf-8"
+ )
monkeypatch.setattr("sys.argv", ["tvm-ffi-stubgen", "--target", "rust",
"--check", str(root)])
assert stub_cli.__main__() == 2
@@ -1997,7 +1997,7 @@ def test_roll_out_matches_the_prefix_exactly(
src.write_text(
"\n".join(
[
- f"{C.RUST_SYNTAX.directive('prefix')} a.",
+ f"{C.RUST_SYNTAX.directive('file')} mod(prefix = a.)",
f"{begin} import-section",
end,
"mod tail {}",
@@ -2010,7 +2010,7 @@ def test_roll_out_matches_the_prefix_exactly(
assert info is not None
assert stub_cli._roll_out_prefixes([info], RUST) == 0
assert [line for block in info.code_blocks for line in block.lines] == [
- f"{C.RUST_SYNTAX.directive('prefix')} a.",
+ f"{C.RUST_SYNTAX.directive('file')} mod(prefix = a.)",
f"{begin} import-section",
end,
"",
@@ -2023,7 +2023,7 @@ def test_roll_out_matches_the_prefix_exactly(
]
# Every `ffi.*` key lives in the crate and is never rolled out,
`ffi.Array` included: a
# generated non-generic `Array` would shadow the crate's `Array<T>`.
- src.write_text(f"{C.RUST_SYNTAX.directive('prefix')} ffi\n",
encoding="utf-8")
+ src.write_text(f"{C.RUST_SYNTAX.directive('file')} mod(prefix = ffi)\n",
encoding="utf-8")
info = FileInfo.from_file(src)
assert info is not None
assert stub_cli._roll_out_prefixes([info], RUST) == 0
@@ -2035,13 +2035,15 @@ def test_prefix_survives_init(tmp_path: Path,
monkeypatch: pytest.MonkeyPatch) -
root = _crate(tmp_path)
(root / "testing").mkdir()
mod_rs = root / "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")
+ skip = f"{C.RUST_SYNTAX.directive('object')}
testing.TestCxxClassDerivedDerived(skip)"
+ mod_rs.write_text(
+ f"{C.RUST_SYNTAX.directive('file')} mod(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(root)])
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 text.startswith(f"{C.RUST_SYNTAX.directive('file')} mod(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(root)])
@@ -2061,7 +2063,13 @@ def test_prefix_removes_blocks_of_unregistered_types(
gone = f"{C.RUST_SYNTAX.begin} object/testing.NoLongerRegistered"
src.write_text(
"\n".join(
- [f"{C.RUST_SYNTAX.directive('prefix')} testing", "", gone,
C.RUST_SYNTAX.end, ""]
+ [
+ f"{C.RUST_SYNTAX.directive('file')} mod(prefix = testing)",
+ "",
+ gone,
+ C.RUST_SYNTAX.end,
+ "",
+ ]
),
encoding="utf-8",
)
