This is an automated email from the ASF dual-hosted git repository.
junrushao 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 ea02e64 feat: stub generation (#101)
ea02e64 is described below
commit ea02e646aea916ad18755cd6990b22ace8afd4b2
Author: Junru Shao <[email protected]>
AuthorDate: Sun Oct 12 10:56:29 2025 -0700
feat: stub generation (#101)
This PR introduces `tvm-ffi-stubgen`, a small CLI tool that generates
in-place, static typing stubs for Python modules that integrate with the
TVM FFI registry. This gives static type checkers, e.g. mypy, rich and
precise type signatures for global functions, object fields and methods.
## Demo
<img width="629" height="881" alt="image"
src="https://github.com/user-attachments/assets/8a3ae92c-78f0-4ca2-b3b6-25e20cc0dac3";
/>
<img width="530" height="341" alt="image"
src="https://github.com/user-attachments/assets/a6d59636-c890-40b5-a679-a7cba822da0d";
/>
## What it does
- Scans .py/.pyi files for special begin/end markers and fills those
blocks with generated code guarded by `if TYPE_CHECKING: ...`
- It honors a skip-file directive to opt out of processing a file
entirely.
- Emits types signatures from TVM FFI's registry, this includes:
* global functions
* object fields
* object methods
- It supports per-block type remapping via `ty_map` lines, e.g. `list ->
Sequence`, `dict -> Mapping`
- Prints a unified diff for each change and writes back in place
deterministically.
## Block directives
- Global functions
* Begin: `# tvm-ffi-stubgen(begin): global/<registry-prefix>`
* End: `# tvm-ffi-stubgen(end)`
* Example:
```python
from typing import TYPE_CHECKING
# tvm-ffi-stubgen(begin): global/ffi
# tvm-ffi-stubgen(end)
```
- Object types
* Begin: `# tvm-ffi-stubgen(begin): object/<type_key>`
* Inside: optional ty_map hints `# tvm-ffi-stubgen(ty_map): A.B -> C`
* End: `# tvm-ffi-stubgen(end)`
* Example:
```python
@register_object("testing.SchemaAllTypes")
class _SchemaAllTypes:
# tvm-ffi-stubgen(begin): object/testing.SchemaAllTypes
# tvm-ffi-stubgen(ty_map): testing.SchemaAllTypes -> _SchemaAllTypes
# tvm-ffi-stubgen(end)
```
- Skip a file
* Anywhere in the file: `# tvm-ffi-stubgen(skip-file)`
## Indentation and output
- Indentation on the begin line is preserved; generated lines are
additionally indented by --indent spaces (default 4).
- If no functions or members are found for a block, nothing is generated
and the block remains unchanged.
- Only .py and .pyi files are modified; directories are scanned
recursively.
## CLI usage
Single file:
```
tvm-ffi-stubgen python/tvm_ffi/_ffi_api.py
```
Recursively scan directories:
```
tvm-ffi-stubgen python/tvm_ffi examples/packaging/python/my_ffi_extension
```
Preload shared libraries (when metadata is provided by native
extensions):
```
tvm-ffi-stubgen --dlls build/libtvm_runtime.so build/libmy_ext.so
my_pkg/_ffi_api.py
```
## Examples in the repo
- Global stubs: python/tvm_ffi/_ffi_api.py
- Object stubs: python/tvm_ffi/testing.py
- Minimal packaged example:
examples/packaging/python/my_ffi_extension/_ffi_api.py:1
## Notes
- Requires Python to import tvm_ffi and for the process to have access
to the TVM runtime and any extension libraries that register the
functions/types being stubbed. Use --dlls to preload them when needed.
- Generated content is deterministic and stable across runs; a second
invocation should leave files unchanged when inputs haven’t changed.
---
.../packaging/python/my_ffi_extension/__init__.py | 2 +-
.../packaging/python/my_ffi_extension/_ffi_api.py | 10 +
pyproject.toml | 1 +
python/tvm_ffi/_ffi_api.py | 51 ++
python/tvm_ffi/_ffi_api.pyi | 43 --
python/tvm_ffi/access_path.py | 83 ++--
python/tvm_ffi/module.py | 16 +-
python/tvm_ffi/{_ffi_api.py => stub/__init__.py} | 6 +-
python/tvm_ffi/stub/stubgen.py | 528 +++++++++++++++++++++
python/tvm_ffi/testing.py | 75 ++-
tests/python/test_object.py | 4 +-
tests/python/test_stubgen.py | 182 +++++++
12 files changed, 905 insertions(+), 96 deletions(-)
diff --git a/examples/packaging/python/my_ffi_extension/__init__.py
b/examples/packaging/python/my_ffi_extension/__init__.py
index 0c2b0fd..ae4abfd 100644
--- a/examples/packaging/python/my_ffi_extension/__init__.py
+++ b/examples/packaging/python/my_ffi_extension/__init__.py
@@ -51,4 +51,4 @@ def raise_error(msg: str) -> None:
The error raised by the function.
"""
- return _ffi_api.raise_error(msg) # type: ignore[attr-defined]
+ return _ffi_api.raise_error(msg)
diff --git a/examples/packaging/python/my_ffi_extension/_ffi_api.py
b/examples/packaging/python/my_ffi_extension/_ffi_api.py
index edc7677..f57eb2a 100644
--- a/examples/packaging/python/my_ffi_extension/_ffi_api.py
+++ b/examples/packaging/python/my_ffi_extension/_ffi_api.py
@@ -14,6 +14,8 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations.
+from typing import TYPE_CHECKING
+
import tvm_ffi
# make sure lib is loaded first
@@ -22,3 +24,11 @@ from .base import _LIB # noqa: F401
# this is a short cut to register all the global functions
# prefixed by `my_ffi_extension.` to this module
tvm_ffi.init_ffi_api("my_ffi_extension", __name__)
+
+
+# tvm-ffi-stubgen(begin): global/my_ffi_extension
+if TYPE_CHECKING:
+ # fmt: off
+ def raise_error(_0: str, /) -> None: ...
+ # fmt: on
+# tvm-ffi-stubgen(end)
diff --git a/pyproject.toml b/pyproject.toml
index a0a00c9..d09eac5 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -73,6 +73,7 @@ docs = [
[project.scripts]
tvm-ffi-config = "tvm_ffi.config:__main__"
+tvm-ffi-stubgen = "tvm_ffi.stub.stubgen:__main__"
[build-system]
requires = ["scikit-build-core>=0.10.0", "cython"]
diff --git a/python/tvm_ffi/_ffi_api.py b/python/tvm_ffi/_ffi_api.py
index f9314dd..5a957d6 100644
--- a/python/tvm_ffi/_ffi_api.py
+++ b/python/tvm_ffi/_ffi_api.py
@@ -16,6 +16,57 @@
# under the License.
"""FFI API."""
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Callable
+
from . import registry
+if TYPE_CHECKING:
+ from collections.abc import Mapping, Sequence
+
+ from tvm_ffi import Module
+ from tvm_ffi.access_path import AccessPath
+
+
+# tvm-ffi-stubgen(begin): global/ffi
+if TYPE_CHECKING:
+ # fmt: off
+ def Array(*args: Any) -> Any: ...
+ def ArrayGetItem(_0: Sequence[Any], _1: int, /) -> Any: ...
+ def ArraySize(_0: Sequence[Any], /) -> int: ...
+ def Bytes(_0: bytes, /) -> bytes: ...
+ def FromJSONGraph(_0: Any, /) -> Any: ...
+ def FromJSONGraphString(_0: str, /) -> Any: ...
+ def FunctionListGlobalNamesFunctor() -> Callable[..., Any]: ...
+ def FunctionRemoveGlobal(_0: str, /) -> bool: ...
+ def GetFirstStructuralMismatch(_0: Any, _1: Any, _2: bool, _3: bool, /) ->
tuple[AccessPath, AccessPath] | None: ...
+ def GetGlobalFuncMetadata(_0: str, /) -> str: ...
+ def MakeObjectFromPackedArgs(*args: Any) -> Any: ...
+ def Map(*args: Any) -> Any: ...
+ def MapCount(_0: Mapping[Any, Any], _1: Any, /) -> int: ...
+ def MapForwardIterFunctor(_0: Mapping[Any, Any], /) -> Callable[..., Any]:
...
+ def MapGetItem(_0: Mapping[Any, Any], _1: Any, /) -> Any: ...
+ def MapSize(_0: Mapping[Any, Any], /) -> int: ...
+ def ModuleClearImports(_0: Module, /) -> None: ...
+ def ModuleGetFunction(_0: Module, _1: str, _2: bool, /) -> Callable[...,
Any] | None: ...
+ def ModuleGetFunctionDoc(_0: Module, _1: str, _2: bool, /) -> str | None:
...
+ def ModuleGetFunctionMetadata(_0: Module, _1: str, _2: bool, /) -> str |
None: ...
+ def ModuleGetKind(_0: Module, /) -> str: ...
+ def ModuleGetPropertyMask(_0: Module, /) -> int: ...
+ def ModuleGetWriteFormats(_0: Module, /) -> Sequence[str]: ...
+ def ModuleImplementsFunction(_0: Module, _1: str, _2: bool, /) -> bool: ...
+ def ModuleImportModule(_0: Module, _1: Module, /) -> None: ...
+ def ModuleInspectSource(_0: Module, _1: str, /) -> str: ...
+ def ModuleLoadFromFile(_0: str, /) -> Module: ...
+ def ModuleWriteToFile(_0: Module, _1: str, _2: str, /) -> None: ...
+ def Shape(*args: Any) -> Any: ...
+ def String(_0: str, /) -> str: ...
+ def StructuralHash(_0: Any, _1: bool, _2: bool, /) -> int: ...
+ def SystemLib(*args: Any) -> Any: ...
+ def ToJSONGraph(_0: Any, _1: Any, /) -> Any: ...
+ def ToJSONGraphString(_0: Any, _1: Any, /) -> str: ...
+ # fmt: on
+# tvm-ffi-stubgen(end)
+
registry.init_ffi_api("ffi", __name__)
diff --git a/python/tvm_ffi/_ffi_api.pyi b/python/tvm_ffi/_ffi_api.pyi
deleted file mode 100644
index 95059e5..0000000
--- a/python/tvm_ffi/_ffi_api.pyi
+++ /dev/null
@@ -1,43 +0,0 @@
-# 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.
-"""FFI API."""
-
-from typing import Any
-
-def ModuleGetKind(*args: Any) -> Any: ...
-def ModuleImplementsFunction(*args: Any) -> Any: ...
-def ModuleGetFunction(*args: Any) -> Any: ...
-def ModuleImportModule(*args: Any) -> Any: ...
-def ModuleInspectSource(*args: Any) -> Any: ...
-def ModuleGetWriteFormats(*args: Any) -> Any: ...
-def ModuleGetPropertyMask(*args: Any) -> Any: ...
-def ModuleClearImports(*args: Any) -> Any: ...
-def ModuleWriteToFile(*args: Any) -> Any: ...
-def ModuleLoadFromFile(*args: Any) -> Any: ...
-def SystemLib(*args: Any) -> Any: ...
-def Array(*args: Any) -> Any: ...
-def ArrayGetItem(*args: Any) -> Any: ...
-def ArraySize(*args: Any) -> Any: ...
-def MapForwardIterFunctor(*args: Any) -> Any: ...
-def Map(*args: Any) -> Any: ...
-def MapGetItem(*args: Any) -> Any: ...
-def MapCount(*args: Any) -> Any: ...
-def MapSize(*args: Any) -> Any: ...
-def MakeObjectFromPackedArgs(*args: Any) -> Any: ...
-def ToJSONGraphString(*args: Any) -> Any: ...
-def FromJSONGraphString(*args: Any) -> Any: ...
-def Shape(*args: Any) -> Any: ...
diff --git a/python/tvm_ffi/access_path.py b/python/tvm_ffi/access_path.py
index aa52d58..c906347 100644
--- a/python/tvm_ffi/access_path.py
+++ b/python/tvm_ffi/access_path.py
@@ -17,10 +17,13 @@
# pylint: disable=invalid-name
"""Access path classes."""
+from __future__ import annotations
+
+from collections.abc import Sequence
from enum import IntEnum
-from typing import Any
+from typing import TYPE_CHECKING, Any
-from . import core
+from .core import Object
from .registry import register_object
@@ -36,18 +39,42 @@ class AccessKind(IntEnum):
@register_object("ffi.reflection.AccessStep")
-class AccessStep(core.Object):
+class AccessStep(Object):
"""Access step container."""
- kind: AccessKind
- key: Any
+ # tvm-ffi-stubgen(begin): object/ffi.reflection.AccessStep
+ if TYPE_CHECKING:
+ # fmt: off
+ kind: int
+ key: Any
+ # fmt: on
+ # tvm-ffi-stubgen(end)
@register_object("ffi.reflection.AccessPath")
-class AccessPath(core.Object):
+class AccessPath(Object):
"""Access path container."""
- parent: "AccessPath"
+ # tvm-ffi-stubgen(begin): object/ffi.reflection.AccessPath
+ if TYPE_CHECKING:
+ # fmt: off
+ parent: Object | None
+ step: AccessStep | None
+ depth: int
+ @staticmethod
+ def _root() -> AccessPath: ...
+ def _extend(_0: AccessPath, _1: AccessStep, /) -> AccessPath: ...
+ def _attr(_0: AccessPath, _1: str, /) -> AccessPath: ...
+ def _array_item(_0: AccessPath, _1: int, /) -> AccessPath: ...
+ def _map_item(_0: AccessPath, _1: Any, /) -> AccessPath: ...
+ def _attr_missing(_0: AccessPath, _1: str, /) -> AccessPath: ...
+ def _array_item_missing(_0: AccessPath, _1: int, /) -> AccessPath: ...
+ def _map_item_missing(_0: AccessPath, _1: Any, /) -> AccessPath: ...
+ def _is_prefix_of(_0: AccessPath, _1: AccessPath, /) -> bool: ...
+ def _to_steps(_0: AccessPath, /) -> Sequence[AccessStep]: ...
+ def _path_equal(_0: AccessPath, _1: AccessPath, /) -> bool: ...
+ # fmt: on
+ # tvm-ffi-stubgen(end)
def __init__(self) -> None:
"""Disallow direct construction; use `AccessPath.root()` instead."""
@@ -58,23 +85,23 @@ class AccessPath(core.Object):
)
@staticmethod
- def root() -> "AccessPath":
+ def root() -> AccessPath:
"""Create a root access path."""
- return AccessPath._root() # type: ignore[attr-defined]
+ return AccessPath._root()
def __eq__(self, other: Any) -> bool:
"""Return whether two access paths are equal."""
if not isinstance(other, AccessPath):
return False
- return self._path_equal(other) # type: ignore[attr-defined]
+ return self._path_equal(other)
def __ne__(self, other: Any) -> bool:
"""Return whether two access paths are not equal."""
if not isinstance(other, AccessPath):
return True
- return not self._path_equal(other) # type: ignore[attr-defined]
+ return not self._path_equal(other)
- def is_prefix_of(self, other: "AccessPath") -> bool:
+ def is_prefix_of(self, other: AccessPath) -> bool:
"""Check if this access path is a prefix of another access path.
Parameters
@@ -88,9 +115,9 @@ class AccessPath(core.Object):
True if this access path is a prefix of the other access path,
False otherwise
"""
- return self._is_prefix_of(other) # type: ignore[attr-defined]
+ return self._is_prefix_of(other)
- def attr(self, attr_key: str) -> "AccessPath":
+ def attr(self, attr_key: str) -> AccessPath:
"""Create an access path to the attribute of the current object.
Parameters
@@ -104,9 +131,9 @@ class AccessPath(core.Object):
The extended access path
"""
- return self._attr(attr_key) # type: ignore[attr-defined]
+ return self._attr(attr_key)
- def attr_missing(self, attr_key: str) -> "AccessPath":
+ def attr_missing(self, attr_key: str) -> AccessPath:
"""Create an access path that indicate an attribute is missing.
Parameters
@@ -120,9 +147,9 @@ class AccessPath(core.Object):
The extended access path
"""
- return self._attr_missing(attr_key) # type: ignore[attr-defined]
+ return self._attr_missing(attr_key)
- def array_item(self, index: int) -> "AccessPath":
+ def array_item(self, index: int) -> AccessPath:
"""Create an access path to the item of the current array.
Parameters
@@ -136,9 +163,9 @@ class AccessPath(core.Object):
The extended access path
"""
- return self._array_item(index) # type: ignore[attr-defined]
+ return self._array_item(index)
- def array_item_missing(self, index: int) -> "AccessPath":
+ def array_item_missing(self, index: int) -> AccessPath:
"""Create an access path that indicate an array item is missing.
Parameters
@@ -152,9 +179,9 @@ class AccessPath(core.Object):
The extended access path
"""
- return self._array_item_missing(index) # type: ignore[attr-defined]
+ return self._array_item_missing(index)
- def map_item(self, key: Any) -> "AccessPath":
+ def map_item(self, key: Any) -> AccessPath:
"""Create an access path to the item of the current map.
Parameters
@@ -168,9 +195,9 @@ class AccessPath(core.Object):
The extended access path
"""
- return self._map_item(key) # type: ignore[attr-defined]
+ return self._map_item(key)
- def map_item_missing(self, key: Any) -> "AccessPath":
+ def map_item_missing(self, key: Any) -> AccessPath:
"""Create an access path that indicate a map item is missing.
Parameters
@@ -184,9 +211,9 @@ class AccessPath(core.Object):
The extended access path
"""
- return self._map_item_missing(key) # type: ignore[attr-defined]
+ return self._map_item_missing(key)
- def to_steps(self) -> list["AccessStep"]:
+ def to_steps(self) -> Sequence[AccessStep]:
"""Convert the access path to a list of access steps.
Returns
@@ -195,6 +222,6 @@ class AccessPath(core.Object):
The list of access steps
"""
- return self._to_steps() # type: ignore[attr-defined]
+ return self._to_steps()
- __hash__ = core.Object.__hash__
+ __hash__ = Object.__hash__
diff --git a/python/tvm_ffi/module.py b/python/tvm_ffi/module.py
index 9503900..9db31a1 100644
--- a/python/tvm_ffi/module.py
+++ b/python/tvm_ffi/module.py
@@ -17,8 +17,9 @@
"""Module related objects and functions."""
# pylint: disable=invalid-name
+from collections.abc import Sequence
from enum import IntEnum
-from typing import Any
+from typing import TYPE_CHECKING, Any, ClassVar, cast
from . import _ffi_api, core
from .registry import register_object
@@ -55,8 +56,14 @@ class Module(core.Object):
"""
- # constant for entry function name
- entry_name = "main"
+ # tvm-ffi-stubgen(begin): object/ffi.Module
+ if TYPE_CHECKING:
+ # fmt: off
+ imports_: Sequence[Any]
+ # fmt: on
+ # tvm-ffi-stubgen(end)
+
+ entry_name: ClassVar[str] = "main" # constant for entry function name
@property
def kind(self) -> str:
@@ -129,6 +136,7 @@ class Module(core.Object):
"""
func = _ffi_api.ModuleGetFunction(self, name, query_imports)
+ func = cast(core.Function, func)
if func is None:
raise AttributeError(f"Module has no function '{name}'")
return func
@@ -171,7 +179,7 @@ class Module(core.Object):
"""
return _ffi_api.ModuleInspectSource(self, fmt)
- def get_write_formats(self) -> list[str]:
+ def get_write_formats(self) -> Sequence[str]:
"""Get the format of the module."""
return _ffi_api.ModuleGetWriteFormats(self)
diff --git a/python/tvm_ffi/_ffi_api.py b/python/tvm_ffi/stub/__init__.py
similarity index 90%
copy from python/tvm_ffi/_ffi_api.py
copy to python/tvm_ffi/stub/__init__.py
index f9314dd..f512cd8 100644
--- a/python/tvm_ffi/_ffi_api.py
+++ b/python/tvm_ffi/stub/__init__.py
@@ -14,8 +14,4 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""FFI API."""
-
-from . import registry
-
-registry.init_ffi_api("ffi", __name__)
+"""A series of stub generator tools."""
diff --git a/python/tvm_ffi/stub/stubgen.py b/python/tvm_ffi/stub/stubgen.py
new file mode 100644
index 0000000..d796eb4
--- /dev/null
+++ b/python/tvm_ffi/stub/stubgen.py
@@ -0,0 +1,528 @@
+# 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.
+"""TVM-FFI Stub Generator (``tvm-ffi-stubgen``).
+
+Overview
+--------
+This module powers the ``tvm-ffi-stubgen`` command line tool which generates
+in-place type stubs for Python modules that integrate with the TVM FFI
+ecosystem. It scans ``.py``/``.pyi`` files for special comment markers and
+fills the enclosed blocks with precise, static type annotations derived from
+runtime metadata exposed by TVM FFI.
+
+Why you might use this
+----------------------
+- You author a Python module that binds to C++/C via TVM FFI and want
+ high-quality type hints for functions, objects, and methods.
+- You maintain a downstream extension that registers global functions or
+ FFI object types and want your Python API surface to be type-checker
+ friendly without manually writing stubs.
+
+How it works (in one sentence)
+------------------------------
+``tvm-ffi-stubgen`` replaces the content between special ``tvm-ffi-stubgen``
markers
+with generated code guarded by ``if TYPE_CHECKING: ...`` so that the runtime
+behavior is unchanged while static analyzers get rich types.
+
+Stub block markers
+------------------
+Insert one of the following begin/end markers in your source, then run
+``tvm-ffi-stubgen``. Indentation on the ``begin`` line is preserved; generated
+content is additionally indented by ``--indent`` spaces (default: 4).
+
+1) Global function stubs
+
+ Mark all global functions whose names start with a registry prefix
+ (e.g. ``ffi`` or ``my_ffi_extension``):
+
+ .. code-block:: python
+
+ from typing import TYPE_CHECKING
+
+ # tvm-ffi-stubgen(begin): global/ffi
+ if TYPE_CHECKING:
+ # fmt: off
+ # (generated by tvm-ffi-stubgen)
+ # fmt: on
+ # tvm-ffi-stubgen(end)
+
+ ``tvm-ffi-stubgen`` expands this with function signatures discovered via
the
+ TVM FFI global function registry.
+
+2) Object type stubs
+
+ Mark fields and methods for a registered FFI object type using its
+ ``type_key`` (the key passed to ``@register_object``):
+
+ .. code-block:: python
+
+ @register_object("testing.SchemaAllTypes")
+ class _SchemaAllTypes:
+ # tvm-ffi-stubgen(begin): object/testing.SchemaAllTypes
+ # tvm-ffi-stubgen(ty_map): testing.SchemaAllTypes ->
_SchemaAllTypes
+ if TYPE_CHECKING:
+ # fmt: off
+ # (generated by tvm-ffi-stubgen)
+ # fmt: on
+ # tvm-ffi-stubgen(end)
+
+ ``tvm-ffi-stubgen`` expands this with annotated attributes and method stub
+ signatures. The special C FFI initializer ``__ffi_init__`` is exposed as
+ ``__c_ffi_init__`` to avoid interfering with your Python ``__init__``.
+
+3) Skip whole file
+
+ If a source file should never be modified by the stub generator, add the
+ following directive anywhere in the file:
+
+ .. code-block:: python
+
+ # tvm-ffi-stubgen(skip-file)
+
+ When present, ``tvm-ffi-stubgen`` skips processing this file entirely. This
+ is useful for files that are generated by other tooling or vendored.
+
+Optional type mapping lines
+---------------------------
+Inside a stub block you may add mapping hints to rename fully-qualified type
+names to simpler aliases in the generated output:
+
+.. code-block:: python
+
+ # tvm-ffi-stubgen(ty_map): A.B.C -> C
+ # tvm-ffi-stubgen(ty_map): list -> Sequence
+ # tvm-ffi-stubgen(ty_map): dict -> Mapping
+
+By default, ``list`` is shown as ``Sequence`` and ``dict`` as ``Mapping``.
+If you use names such as ``Sequence``/``Mapping``, ensure they are available
+to type checkers in your module, for example:
+
+.. code-block:: python
+
+ from typing import TYPE_CHECKING
+ if TYPE_CHECKING:
+ from collections.abc import Mapping, Sequence
+
+Runtime requirements
+--------------------
+- Python must be able to import ``tvm_ffi``.
+- The process needs access to the TVM runtime and any extension libraries that
+ provide the global functions or object types you want to stub. Use the
+ ``--dlls`` option to preload shared libraries when necessary.
+
+What files are modified
+-----------------------
+Only files with extensions ``.py`` and ``.pyi`` are scanned. Files are updated
+in place. A colored unified diff is printed for each change.
+
+CLI quick start
+---------------
+
+.. code-block:: bash
+
+ # Generate stubs for a single file
+ tvm-ffi-stubgen python/tvm_ffi/_ffi_api.py
+
+ # Recursively scan directories for tvm-ffi-stubgen blocks
+ tvm-ffi-stubgen python/tvm_ffi examples/packaging/python/my_ffi_extension
+
+ # Preload TVM runtime and your extension library before generation
+ tvm-ffi-stubgen \
+ --dlls build/libtvm_runtime.dylib build/libmy_ext.dylib \
+ python/tvm_ffi/_ffi_api.py
+
+Exit status
+-----------
+Returns 0 on success and 1 if any file fails to process.
+
+"""
+
+from __future__ import annotations
+
+import argparse
+import ctypes
+import dataclasses
+import difflib
+import logging
+import sys
+from io import StringIO
+from pathlib import Path
+from typing import Callable
+
+from tvm_ffi.core import TypeSchema,
_lookup_or_register_type_info_from_type_key
+from tvm_ffi.registry import get_global_func_metadata, list_global_func_names
+
+DEFAULT_SOURCE_EXTS = {".py", ".pyi"}
+STUB_BEGIN = "# tvm-ffi-stubgen(begin):"
+STUB_END = "# tvm-ffi-stubgen(end)"
+STUB_TY_MAP = "# tvm-ffi-stubgen(ty_map):"
+STUB_SKIP_FILE = "# tvm-ffi-stubgen(skip-file)"
+
+TERM_RESET = "\033[0m"
+TERM_BOLD = "\033[1m"
+TERM_RED = "\033[31m"
+TERM_GREEN = "\033[32m"
+TERM_YELLOW = "\033[33m"
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+
[email protected]
+class Options:
+ """Command line options for stub generation."""
+
+ dlls: list[str] = dataclasses.field(default_factory=list)
+ indent: int = 4
+ files: list[str] = dataclasses.field(default_factory=list)
+ suppress_print: bool = False
+
+
[email protected]
+class StubConfig:
+ """Configuration of a stub block."""
+
+ name: str
+ indent: int
+ lineno: int
+ ty_map: dict[str, str] = dataclasses.field(
+ default_factory=lambda: dict(
+ {
+ "list": "Sequence",
+ "dict": "Mapping",
+ }
+ )
+ )
+
+
+def _as_func_signature(
+ schema: TypeSchema,
+ func_name: str,
+ ty_map: Callable[[str], str],
+) -> str:
+ buf = StringIO()
+ buf.write(f"def {func_name}(")
+ if schema.origin != "Callable":
+ raise ValueError(f"Expected Callable type schema, but got: {schema}")
+ if not schema.args:
+ buf.write("*args: Any) -> Any:")
+ return buf.getvalue()
+ arg_ret = schema.args[0]
+ arg_args = schema.args[1:]
+ for i, arg in enumerate(arg_args):
+ buf.write(f"_{i}: ")
+ buf.write(arg.repr(ty_map))
+ buf.write(", ")
+ if arg_args:
+ buf.write("/")
+ buf.write(") -> ")
+ buf.write(arg_ret.repr(ty_map))
+ buf.write(":")
+ return buf.getvalue()
+
+
+def _filter_files(paths: list[Path]) -> list[Path]:
+ results: list[Path] = []
+ for p in paths:
+ if not p.exists():
+ raise FileNotFoundError(f"Path does not exist: {p}")
+ if p.is_dir():
+ for f in p.rglob("*"):
+ if f.is_file() and f.suffix.lower() in DEFAULT_SOURCE_EXTS:
+ results.append(f.resolve())
+ continue
+ f = p.resolve()
+ if f.is_file() and f.suffix.lower() in DEFAULT_SOURCE_EXTS:
+ results.append(f)
+ # Deterministic order
+ return sorted(set(results))
+
+
+def _make_type_map(name_map: dict[str, str]) -> Callable[[str], str]:
+ def map_type(name: str) -> str:
+ if (ret := name_map.get(name)) is not None:
+ return ret
+ return name.rsplit(".", 1)[-1]
+
+ return map_type
+
+
+def _generate_global(
+ stub: StubConfig,
+ global_func_tab: dict[str, list[str]],
+ opt: Options,
+) -> list[str]:
+ assert stub.name.startswith("global/")
+ prefix = stub.name[len("global/") :].strip()
+ ty_map = _make_type_map(stub.ty_map)
+ indent = " " * (stub.indent + opt.indent)
+ results: list[str] = [
+ " " * stub.indent + "if TYPE_CHECKING:",
+ f"{indent}# fmt: off",
+ ]
+ for name in global_func_tab.get(prefix, []):
+ schema_str =
get_global_func_metadata(f"{prefix}.{name}")["type_schema"]
+ schema = TypeSchema.from_json_str(schema_str)
+ sig = _as_func_signature(schema, name, ty_map=ty_map)
+ func = f"{indent}{sig} ..."
+ results.append(func)
+ if len(results) > 2:
+ results.append(f"{indent}# fmt: on")
+ else:
+ results = []
+ return results
+
+
+def _show_diff(old: list[str], new: list[str]) -> None:
+ for line in difflib.unified_diff(old, new, lineterm=""):
+ # Skip placeholder headers when fromfile/tofile are unspecified
+ if line.startswith("---") or line.startswith("+++"):
+ continue
+ if line.startswith("-") and not line.startswith("---"):
+ print(f"{TERM_RED}{line}{TERM_RESET}") # Red for removals
+ elif line.startswith("+") and not line.startswith("+++"):
+ print(f"{TERM_GREEN}{line}{TERM_RESET}") # Green for additions
+ elif line.startswith("?"):
+ print(f"{TERM_YELLOW}{line}{TERM_RESET}") # Yellow for hints
+ else:
+ print(line)
+
+
+def _generate_object(
+ stub: StubConfig,
+ opt: Options,
+) -> list[str]:
+ assert stub.name.startswith("object/")
+ type_key = stub.name[len("object/") :].strip()
+ ty_map = _make_type_map(stub.ty_map)
+ indent = " " * (stub.indent + opt.indent)
+ results: list[str] = [
+ " " * stub.indent + "if TYPE_CHECKING:",
+ f"{indent}# fmt: off",
+ ]
+
+ type_info = _lookup_or_register_type_info_from_type_key(type_key)
+ for field in type_info.fields:
+ schema = TypeSchema.from_json_str(field.metadata["type_schema"])
+ schema_str = schema.repr(ty_map=ty_map)
+ results.append(f"{indent}{field.name}: {schema_str}")
+ for method in type_info.methods:
+ name = method.name
+ if name == "__ffi_init__":
+ name = "__c_ffi_init__"
+ schema = TypeSchema.from_json_str(method.metadata["type_schema"])
+ schema_str = _as_func_signature(schema, name, ty_map=ty_map)
+ if method.is_static:
+ results.append(f"{indent}@staticmethod")
+ results.append(f"{indent}{schema_str} ...")
+ if len(results) > 2:
+ results.append(f"{indent}# fmt: on")
+ else:
+ results = []
+ return results
+
+
+def _has_skip_file_marker(lines: list[str]) -> bool:
+ for raw in lines:
+ if raw.strip().startswith(STUB_SKIP_FILE):
+ return True
+ return False
+
+
+def _main( # noqa: PLR0912, PLR0915
+ file: Path,
+ opt: Options,
+ global_func_tab: dict[str, list[str]] | None = None,
+) -> None:
+ assert file.is_file(), f"Expected a file, but got: {file}"
+
+ lines_now = file.read_text(encoding="utf-8").splitlines()
+
+ # directive(skip-file): skip processing this file entirely if present.
+ if _has_skip_file_marker(lines_now):
+ if not opt.suppress_print:
+ print(f"{TERM_YELLOW}[Skipped] {file}{TERM_RESET}")
+ return
+
+ if global_func_tab is None:
+ global_func_tab = _compute_global_func_tab()
+
+ lines_new: list[str] = []
+ stub: StubConfig | None = None
+ skipped: bool = True
+ for lineno, line in enumerate(lines_now, 1):
+ clean_line = line.strip()
+ if clean_line.startswith(STUB_BEGIN):
+ if stub is not None:
+ raise ValueError(f"Nested stub not permitted, but found at
{file}:{lineno}")
+ stub = StubConfig(
+ name=clean_line[len(STUB_BEGIN) :].strip(),
+ indent=len(line) - len(clean_line),
+ lineno=lineno,
+ )
+ skipped = False
+ lines_new.append(line)
+ elif clean_line.startswith(STUB_END):
+ if stub is None:
+ raise ValueError(f"Unmatched stub end found at
{file}:{lineno}")
+ if stub.name.startswith("global/"):
+ lines_new.extend(_generate_global(stub, global_func_tab, opt))
+ elif stub.name.startswith("object/"):
+ lines_new.extend(_generate_object(stub, opt))
+ else:
+ raise ValueError(f"Unknown stub type `{stub.name}` at
{file}:{stub.lineno}")
+ stub = None
+ lines_new.append(line)
+ elif clean_line.startswith(STUB_TY_MAP):
+ if stub is None:
+ raise ValueError(f"Stub ty_map outside stub block at
{file}:{lineno}")
+ ty_map = clean_line[len(STUB_TY_MAP) :].strip()
+ try:
+ lhs, rhs = ty_map.split("->")
+ except ValueError as e:
+ raise ValueError(
+ f"Invalid ty_map format at {file}:{lineno}. Example: `A.B
-> C`"
+ ) from e
+ lhs = lhs.strip()
+ rhs = rhs.strip()
+ stub.ty_map[lhs] = rhs
+ lines_new.append(line)
+ elif stub is None:
+ lines_new.append(line)
+ if stub is not None:
+ raise ValueError(f"Unclosed stub block at end of file: {file}")
+ if not skipped:
+ if lines_now != lines_new:
+ if not opt.suppress_print:
+ print(f"{TERM_GREEN}[Updated] {file}{TERM_RESET}")
+ _show_diff(lines_now, lines_new)
+ file.write_text("\n".join(lines_new) + "\n", encoding="utf-8")
+ elif not opt.suppress_print:
+ print(f"{TERM_BOLD}[Unchanged] {file}{TERM_RESET}")
+
+
+def _compute_global_func_tab() -> dict[str, list[str]]:
+ # Build global function table only if we are going to process blocks.
+ global_func_tab: dict[str, list[str]] = {}
+ for name in list_global_func_names():
+ prefix, suffix = name.rsplit(".", 1)
+ global_func_tab.setdefault(prefix, []).append(suffix)
+ # Ensure stable ordering for deterministic output.
+ for k in list(global_func_tab.keys()):
+ global_func_tab[k].sort()
+ return global_func_tab
+
+
+def __main__() -> int:
+ """Command line entry point for ``tvm-ffi-stubgen``.
+
+ This generates in-place type stubs inside special ``tvm-ffi-stubgen``
blocks
+ in the given files or directories. See the module docstring for an
+ overview and examples of the block syntax.
+ """
+
+ class HelpFormatter(argparse.ArgumentDefaultsHelpFormatter,
argparse.RawTextHelpFormatter):
+ pass
+
+ parser = argparse.ArgumentParser(
+ prog="tvm-ffi-stubgen",
+ description=(
+ "Generate in-place type stubs for TVM FFI.\n\n"
+ "It scans .py/.pyi files for tvm-ffi-stubgen blocks and fills them
with\n"
+ "TYPE_CHECKING-only annotations derived from TVM runtime metadata."
+ ),
+ formatter_class=HelpFormatter,
+ epilog=(
+ "Examples:\n"
+ " # Single file\n"
+ " tvm-ffi-stubgen python/tvm_ffi/_ffi_api.py\n\n"
+ " # Recursively scan directories\n"
+ " tvm-ffi-stubgen python/tvm_ffi
examples/packaging/python/my_ffi_extension\n\n"
+ " # Preload TVM runtime / extension libraries\n"
+ " tvm-ffi-stubgen --dlls build/libtvm_runtime.so
build/libmy_ext.so my_pkg/_ffi_api.py\n\n"
+ "Stub block syntax (placed in your source):\n"
+ " # tvm-ffi-stubgen(begin): global/<registry-prefix>\n"
+ " ... generated function stubs ...\n"
+ " # tvm-ffi-stubgen(end)\n\n"
+ " # tvm-ffi-stubgen(begin): object/<type_key>\n"
+ " # tvm-ffi-stubgen(ty_map): list -> Sequence\n"
+ " # tvm-ffi-stubgen(ty_map): dict -> Mapping\n"
+ " ... generated fields and methods ...\n"
+ " # tvm-ffi-stubgen(end)\n\n"
+ " # Skip a file entirely\n"
+ " # tvm-ffi-stubgen(skip-file)\n\n"
+ "Tips:\n"
+ " - Only .py/.pyi files are updated; directories are scanned
recursively.\n"
+ " - Import any aliases you use in ty_map under TYPE_CHECKING,
e.g.\n"
+ " from collections.abc import Mapping, Sequence\n"
+ " - Use --dlls to preload shared libraries when function/type
metadata\n"
+ " is provided by native extensions.\n"
+ ),
+ )
+ parser.add_argument(
+ "--dlls",
+ nargs="*",
+ metavar="LIB",
+ help=(
+ "Shared libraries to preload before generation (e.g. TVM runtime
or "
+ "your extension). This ensures global function and object metadata
"
+ "is available. Accepts multiple paths; platform-specific suffixes "
+ "like .so/.dylib/.dll are supported."
+ ),
+ default=[],
+ )
+ parser.add_argument(
+ "--indent",
+ type=int,
+ default=4,
+ help=(
+ "Extra spaces added inside each generated block, relative to the "
+ "indentation of the corresponding '# tvm-ffi-stubgen(begin):'
line."
+ ),
+ )
+ parser.add_argument(
+ "files",
+ nargs="*",
+ metavar="PATH",
+ help=(
+ "Files or directories to process. Directories are scanned
recursively; "
+ "only .py and .pyi files are modified. Use tvm-ffi-stubgen markers
to "
+ "select where stubs are generated."
+ ),
+ )
+ opt = Options(**vars(parser.parse_args()))
+ if not opt.files:
+ parser.print_help()
+ return 1
+
+ dlls = [ctypes.CDLL(lib) for lib in opt.dlls]
+ global_func_tab = _compute_global_func_tab()
+ rc = 0
+ try:
+ for file in _filter_files([Path(f) for f in opt.files]):
+ try:
+ _main(file, opt, global_func_tab=global_func_tab)
+ except Exception:
+ logger.exception(f"{TERM_RED}[Failed] {file}{TERM_RESET}")
+ rc = 1
+ finally:
+ del dlls
+ return rc
+
+
+if __name__ == "__main__":
+ sys.exit(__main__())
diff --git a/python/tvm_ffi/testing.py b/python/tvm_ffi/testing.py
index 097e298..2ea89ce 100644
--- a/python/tvm_ffi/testing.py
+++ b/python/tvm_ffi/testing.py
@@ -15,33 +15,50 @@
# specific language governing permissions and limitations
# under the License.
"""Testing utilities."""
+# ruff: noqa: D102,D105
from __future__ import annotations
-from typing import Any, ClassVar
+from collections.abc import Mapping, Sequence
+from typing import TYPE_CHECKING, Any, ClassVar
from . import _ffi_api
-from .container import Array, Map
from .core import Object
from .dataclasses import c_class, field
from .registry import get_global_func, register_object
+if TYPE_CHECKING:
+ from tvm_ffi import Device, dtype
+
@register_object("testing.TestObjectBase")
class TestObjectBase(Object):
"""Test object base class."""
- v_i64: int
- v_f64: float
- v_str: str
+ # tvm-ffi-stubgen(begin): object/testing.TestObjectBase
+ if TYPE_CHECKING:
+ # fmt: off
+ v_i64: int
+ v_f64: float
+ v_str: str
+ def add_i64(_0: TestObjectBase, _1: int, /) -> int: ...
+ # fmt: on
+ # tvm-ffi-stubgen(end)
@register_object("testing.TestIntPair")
class TestIntPair(Object):
"""Test Int Pair."""
- a: int
- b: int
+ # tvm-ffi-stubgen(begin): object/testing.TestIntPair
+ if TYPE_CHECKING:
+ # fmt: off
+ a: int
+ b: int
+ @staticmethod
+ def __c_ffi_init__(_0: int, _1: int, /) -> Object: ...
+ # fmt: on
+ # tvm-ffi-stubgen(end)
def __init__(self, a: int, b: int) -> None:
"""Construct the object."""
@@ -52,8 +69,44 @@ class TestIntPair(Object):
class TestObjectDerived(TestObjectBase):
"""Test object derived class."""
- v_map: Map[Any, Any]
- v_array: Array[Any]
+ # tvm-ffi-stubgen(begin): object/testing.TestObjectDerived
+ if TYPE_CHECKING:
+ # fmt: off
+ v_map: Mapping[Any, Any]
+ v_array: Sequence[Any]
+ # fmt: on
+ # tvm-ffi-stubgen(end)
+
+
+@register_object("testing.SchemaAllTypes")
+class _SchemaAllTypes:
+ # tvm-ffi-stubgen(begin): object/testing.SchemaAllTypes
+ # tvm-ffi-stubgen(ty_map): testing.SchemaAllTypes -> _SchemaAllTypes
+ if TYPE_CHECKING:
+ # fmt: off
+ v_bool: bool
+ v_int: int
+ v_float: float
+ v_device: Device
+ v_dtype: dtype
+ v_string: str
+ v_bytes: bytes
+ v_opt_int: int | None
+ v_opt_str: str | None
+ v_arr_int: Sequence[int]
+ v_arr_str: Sequence[str]
+ v_map_str_int: Mapping[str, int]
+ v_map_str_arr_int: Mapping[str, Sequence[int]]
+ v_variant: str | Sequence[int] | Mapping[str, int]
+ v_opt_arr_variant: Sequence[int | str] | None
+ def add_int(_0: _SchemaAllTypes, _1: int, /) -> int: ...
+ def append_int(_0: _SchemaAllTypes, _1: Sequence[int], _2: int, /) ->
Sequence[int]: ...
+ def maybe_concat(_0: _SchemaAllTypes, _1: str | None, _2: str | None,
/) -> str | None: ...
+ def merge_map(_0: _SchemaAllTypes, _1: Mapping[str, Sequence[int]],
_2: Mapping[str, Sequence[int]], /) -> Mapping[str, Sequence[int]]: ...
+ @staticmethod
+ def make_with(_0: int, _1: float, _2: str, /) -> _SchemaAllTypes: ...
+ # fmt: on
+ # tvm-ffi-stubgen(end)
def create_object(type_key: str, **kwargs: Any) -> Object:
@@ -122,7 +175,3 @@ class _TestCxxInitSubset:
required_field: int
optional_field: int = field(init=False)
note: str = field(default_factory=lambda: "py-default", init=False)
-
-
-@register_object("testing.SchemaAllTypes")
-class _SchemaAllTypes: ...
diff --git a/tests/python/test_object.py b/tests/python/test_object.py
index 12dcf3f..98c0bb0 100644
--- a/tests/python/test_object.py
+++ b/tests/python/test_object.py
@@ -90,8 +90,8 @@ def test_derived_object() -> None:
"testing.TestObjectDerived", v_i64=20, v_map=v_map, v_array=v_array
)
assert isinstance(obj0, tvm_ffi.testing.TestObjectDerived)
- assert obj0.v_map.same_as(v_map)
- assert obj0.v_array.same_as(v_array)
+ assert obj0.v_map.same_as(v_map) # type: ignore[attr-defined]
+ assert obj0.v_array.same_as(v_array) # type: ignore[attr-defined]
assert obj0.v_i64 == 20
assert obj0.v_f64 == 10.0
assert obj0.v_str == "hello"
diff --git a/tests/python/test_stubgen.py b/tests/python/test_stubgen.py
new file mode 100644
index 0000000..574d2b3
--- /dev/null
+++ b/tests/python/test_stubgen.py
@@ -0,0 +1,182 @@
+# 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.
+from __future__ import annotations
+
+import re
+from pathlib import Path
+
+import pytest
+from tvm_ffi.stub import stubgen
+
+
+def test_stubgen_skip_file(tmp_path: Path) -> None:
+ p: Path = tmp_path / "dummy.py"
+ src = (
+ "# tvm-ffi-stubgen(skip-file)\n"
+ "from typing import TYPE_CHECKING\n\n"
+ "# tvm-ffi-stubgen(begin): global/ffi\n"
+ "if TYPE_CHECKING:\n"
+ " pass\n"
+ "# tvm-ffi-stubgen(end)\n"
+ )
+ p.write_text(src, encoding="utf-8")
+ # Run the generator; it should skip without trying to query the registry
+ stubgen._main(p, stubgen.Options(indent=4, suppress_print=True))
+ # File must be unchanged
+ assert p.read_text(encoding="utf-8") == src
+
+
+def test_stubgen_global_block_generates_and_indents(tmp_path: Path) -> None:
+ p: Path = tmp_path / "gen_global.py"
+ # Indent begin by 2 spaces; inner indent is begin-indent + opt.indent (3)
+ src = (
+ "from typing import TYPE_CHECKING\n\n"
+ " # tvm-ffi-stubgen(begin): global/ffi\n"
+ " # tvm-ffi-stubgen(end)\n"
+ )
+ p.write_text(src, encoding="utf-8")
+
+ stubgen._main(p, stubgen.Options(indent=3, suppress_print=True))
+
+ out = p.read_text(encoding="utf-8").splitlines()
+
+ # Expect TYPE_CHECKING guard with the same begin indentation
+ assert any(line == " if TYPE_CHECKING:" for line in out)
+ # Expect formatting guards
+ assert any(line == " # fmt: off" for line in out)
+ assert any(line == " # fmt: on" for line in out)
+ # Expect at least one known ffi function signature, e.g. String(...)-> str
+ string_lines = [ln for ln in out if
re.search(r"\bdef\s+String\(.*\)\s*->\s*str:\s*\.\.\.", ln)]
+ assert string_lines, "Expected stub for ffi.String"
+ # Check inner indent equals begin (2) + opt.indent (3) = 5 spaces
+ assert all(ln.startswith(" " * 5) for ln in string_lines)
+
+ # Idempotency: second run should keep file unchanged
+ before = "\n".join(out) + "\n"
+ stubgen._main(p, stubgen.Options(indent=3, suppress_print=True))
+ assert p.read_text(encoding="utf-8") == before
+
+
+def test_stubgen_global_block_no_matches_is_noop(tmp_path: Path) -> None:
+ p: Path = tmp_path / "gen_global_empty.py"
+ src = "# tvm-ffi-stubgen(begin): global/this_prefix_does_not_exist\n#
tvm-ffi-stubgen(end)\n"
+ p.write_text(src, encoding="utf-8")
+ stubgen._main(p, stubgen.Options(indent=4, suppress_print=True))
+ assert p.read_text(encoding="utf-8") == src
+
+
+def test_stubgen_object_block_generates_fields_and_methods(tmp_path: Path) ->
None:
+ # Ensure object type registrations are loaded
+
+ p: Path = tmp_path / "gen_object_pair.py"
+ src = (
+ "class _C:\n"
+ " # tvm-ffi-stubgen(begin): object/testing.TestIntPair\n"
+ " # tvm-ffi-stubgen(end)\n"
+ )
+ p.write_text(src, encoding="utf-8")
+
+ stubgen._main(p, stubgen.Options(indent=4, suppress_print=True))
+ out = p.read_text(encoding="utf-8").splitlines()
+
+ # Fields a and b should be generated
+ assert any(" a: int" == ln for ln in out)
+ assert any(" b: int" == ln for ln in out)
+ # __ffi_init__ should be exposed as __c_ffi_init__ and marked staticmethod
+ init_idx = next(i for i, ln in enumerate(out) if "def __c_ffi_init__(" in
ln)
+ assert out[init_idx - 1].strip() == "@staticmethod"
+
+
+def test_stubgen_object_block_with_ty_map_and_collections(tmp_path: Path) ->
None:
+ # Ensure type info for SchemaAllTypes is available
+ p: Path = tmp_path / "gen_object_schema.py"
+ src = (
+ "# tvm-ffi-stubgen(begin): object/testing.SchemaAllTypes\n"
+ "# tvm-ffi-stubgen(ty_map): testing.SchemaAllTypes ->
_SchemaAllTypes\n"
+ "# tvm-ffi-stubgen(end)\n"
+ )
+ p.write_text(src, encoding="utf-8")
+
+ stubgen._main(p, stubgen.Options(indent=4, suppress_print=True))
+ text = p.read_text(encoding="utf-8")
+
+ # Mapped container aliases should appear
+ assert "Sequence[int]" in text
+ assert "Mapping[str, Sequence[int]]" in text
+ # Method types reflect mapping of the object type
+ assert re.search(r"def\s+add_int\(_0: _SchemaAllTypes, _1: int,
/\)\s*->\s*int:\s*\.\.\.", text)
+ # Static factory returns the mapped type
+ assert re.search(
+
r"@staticmethod\s*\n\s*def\s+make_with\(.*\)\s*->\s*_SchemaAllTypes:\s*\.\.\.",
text
+ )
+
+
+def test_stubgen_errors_for_invalid_directives(tmp_path: Path) -> None:
+ # ty_map outside a block
+ p1 = tmp_path / "invalid_ty_map_outside.py"
+ p1.write_text("# tvm-ffi-stubgen(ty_map): A.B -> C\n", encoding="utf-8")
+ with pytest.raises(ValueError, match="Stub ty_map outside stub block"):
+ stubgen._main(p1, stubgen.Options(suppress_print=True))
+
+ # invalid ty_map format inside a block
+ p2 = tmp_path / "invalid_ty_map_format.py"
+ p2.write_text(
+ (
+ "# tvm-ffi-stubgen(begin): object/testing.TestObjectBase\n"
+ "# tvm-ffi-stubgen(ty_map): not_a_map_line\n"
+ "# tvm-ffi-stubgen(end)\n"
+ ),
+ encoding="utf-8",
+ )
+ with pytest.raises(ValueError, match=r"Invalid ty_map format"):
+ stubgen._main(p2, stubgen.Options(suppress_print=True))
+
+
+def test_stubgen_errors_for_block_structure(tmp_path: Path) -> None:
+ # Nested stub blocks are not allowed
+ p_nested = tmp_path / "nested.py"
+ p_nested.write_text(
+ (
+ "# tvm-ffi-stubgen(begin): global/ffi\n"
+ " # tvm-ffi-stubgen(begin): global/ffi\n"
+ "# tvm-ffi-stubgen(end)\n"
+ ),
+ encoding="utf-8",
+ )
+ with pytest.raises(ValueError, match=r"Nested stub not permitted"):
+ stubgen._main(p_nested, stubgen.Options(suppress_print=True))
+
+ # Unmatched end
+ p_unmatched_end = tmp_path / "unmatched_end.py"
+ p_unmatched_end.write_text("# tvm-ffi-stubgen(end)\n", encoding="utf-8")
+ with pytest.raises(ValueError, match=r"Unmatched stub end"):
+ stubgen._main(p_unmatched_end, stubgen.Options(suppress_print=True))
+
+ # Unknown stub type
+ p_unknown = tmp_path / "unknown.py"
+ p_unknown.write_text(
+ ("# tvm-ffi-stubgen(begin): unknown/foo\n# tvm-ffi-stubgen(end)\n"),
+ encoding="utf-8",
+ )
+ with pytest.raises(ValueError, match=r"Unknown stub type"):
+ stubgen._main(p_unknown, stubgen.Options(suppress_print=True))
+
+ # Unclosed block
+ p_unclosed = tmp_path / "unclosed.py"
+ p_unclosed.write_text("# tvm-ffi-stubgen(begin): global/ffi\n",
encoding="utf-8")
+ with pytest.raises(ValueError, match=r"Unclosed stub block"):
+ stubgen._main(p_unclosed, stubgen.Options(suppress_print=True))
