junrushao opened a new pull request, #101:
URL: https://github.com/apache/tvm-ffi/pull/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.
## 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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
