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 dc0dd2f doc: Improve python packaging doc (#355)
dc0dd2f is described below
commit dc0dd2f6e8367e3bb244e0dc09c6a1a41495499d
Author: Junru Shao <[email protected]>
AuthorDate: Mon Dec 22 06:38:27 2025 -0800
doc: Improve python packaging doc (#355)
Made a few changes:
- Migrate `examples/packaging` -> `examples/python_packaging`
- Add a prerequisite section
- Use `literalinclude` to include code from `examples/python_packaging`
to the tutorial
- Point to Python package docs in `tvm-ffi-stubgen` helper messages
---
.github/workflows/ci_mainline_only.yml | 8 +-
docs/packaging/python_packaging.rst | 104 +++++--------
examples/packaging/src/extension.cc | 131 -----------------
.../{packaging => python_packaging}/CMakeLists.txt | 4 +-
examples/{packaging => python_packaging}/README.md | 18 +--
.../{packaging => python_packaging}/pyproject.toml | 14 +-
.../python/my_ffi_extension/__init__.py | 0
.../python/my_ffi_extension/_ffi_api.py | 6 +-
.../{packaging => python_packaging}/run_example.py | 24 +--
examples/python_packaging/src/extension.cc | 91 ++++++++++++
python/tvm_ffi/stub/cli.py | 163 +--------------------
python/tvm_ffi/stub/consts.py | 1 +
12 files changed, 175 insertions(+), 389 deletions(-)
diff --git a/.github/workflows/ci_mainline_only.yml
b/.github/workflows/ci_mainline_only.yml
index 1912acb..dbf170b 100644
--- a/.github/workflows/ci_mainline_only.yml
+++ b/.github/workflows/ci_mainline_only.yml
@@ -163,25 +163,25 @@ jobs:
cd examples\stable_c_abi
call run_all.bat
- - name: Run example/packaging [posix]
+ - name: Run example/python_packaging [posix]
if: ${{ runner.os != 'Windows' }}
env:
CMAKE_BUILD_PARALLEL_LEVEL: ${{ steps.env_vars.outputs.cpu_count }}
run: |
- pushd examples/packaging
+ pushd examples/python_packaging
# This directory will be auto-generated in `CMakeLists.txt` by
setting `STUB_INIT ON`
rm -rf python/my_ffi_extension
uv pip install --verbose . --no-build-isolation
python run_example.py
popd
- - name: Run example/packaging [windows]
+ - name: Run example/python_packaging [windows]
if: ${{ runner.os == 'Windows' }}
shell: pwsh
env:
CMAKE_BUILD_PARALLEL_LEVEL: ${{ steps.env_vars.outputs.cpu_count }}
run: |
- Set-Location examples/packaging
+ Set-Location examples/python_packaging
Remove-Item -Recurse -Force python/my_ffi_extension
uv pip install --verbose . --no-build-isolation
python run_example.py
diff --git a/docs/packaging/python_packaging.rst
b/docs/packaging/python_packaging.rst
index c768768..3166fdb 100644
--- a/docs/packaging/python_packaging.rst
+++ b/docs/packaging/python_packaging.rst
@@ -27,6 +27,23 @@ internals. We will cover three checkpoints:
- Build Python wheel;
- Automatic Python package generation tools.
+.. note::
+
+ All code used in this guide lives under
+ `examples/python_packaging
<https://github.com/apache/tvm-ffi/tree/main/examples/python_packaging>`_.
+
+.. admonition:: Prerequisite
+ :class: hint
+
+ - Python: 3.9 or newer (for the ``tvm_ffi.config``/``tvm-ffi-config``
helpers)
+ - Compiler: C11-capable toolchain (GCC/Clang/MSVC)
+ - TVM-FFI installed via
+
+ .. code-block:: bash
+
+ pip install --reinstall --upgrade apache-tvm-ffi
+
+
Export C++ to Python
--------------------
@@ -53,13 +70,10 @@ C symbols are easier to call into.
Macro :c:macro:`TVM_FFI_DLL_EXPORT_TYPED_FUNC` exports the function
``AddTwo`` as
a C symbol ``__tvm_ffi_add_two`` inside the shared library.
- .. code-block:: cpp
-
- static int AddTwo(int x) {
- return x + 2;
- }
-
- TVM_FFI_DLL_EXPORT_TYPED_FUNC(add_two, AddTwo);
+ .. literalinclude:: ../../examples/python_packaging/src/extension.cc
+ :language: cpp
+ :start-after: [tvm_ffi_abi.begin]
+ :end-before: [tvm_ffi_abi.end]
.. group-tab:: Python (User)
@@ -97,17 +111,10 @@ It registry handles type translation, error handling, and
metadata.
C++ function ``AddOne`` is registered with name
``my_ffi_extension.add_one``
in the global registry using :cpp:class:`tvm::ffi::reflection::GlobalDef`.
- .. code-block:: cpp
-
- static int AddOne(int x) {
- return x + 1;
- }
-
- TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::GlobalDef()
- .def("my_ffi_extension.add_one", AddOne);
- }
+ .. literalinclude:: ../../examples/python_packaging/src/extension.cc
+ :language: cpp
+ :start-after: [global_function.begin]
+ :end-before: [global_function.end]
.. group-tab:: Python (User)
@@ -166,33 +173,10 @@ makes it easy to expose:
- a constructor, and
- a method ``Sum`` that returns the sum of the two fields.
- .. code-block:: cpp
-
- class IntPairObj : public ffi::Object {
- public:
- int64_t a;
- int64_t b;
- IntPairObj(int64_t a, int64_t b) : a(a), b(b) {}
-
- int64_t Sum() const {
- return a + b;
- }
-
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL(
- /*type_key=*/"my_ffi_extension.IntPair",
- /*class=*/IntPairObj,
- /*parent_class=*/ffi::Object
- );
- };
-
- TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- refl::ObjectDef<IntPairObj>()
- .def(refl::init<int64_t, int64_t>())
- .def_rw("a", &IntPairObj::a, "the first field")
- .def_rw("b", &IntPairObj::b, "the second field")
- .def("sum", &IntPairObj::Sum, "IntPairObj::Sum() method");
- }
+ .. literalinclude:: ../../examples/python_packaging/src/extension.cc
+ :language: cpp
+ :start-after: [object.begin]
+ :end-before: [object.end]
.. group-tab:: Python (User)
@@ -238,15 +222,14 @@ CMake Target
Assume the source tree contains ``src/extension.cc``. Create a
``CMakeLists.txt`` that
creates a shared target ``my_ffi_extension`` and configures it against TVM-FFI.
-.. code-block:: cmake
-
- add_library(my_ffi_extension SHARED src/extension.cc)
- tvm_ffi_configure_target(my_ffi_extension STUB_DIR "./python")
- install(TARGETS my_ffi_extension DESTINATION .)
- tvm_ffi_install(my_ffi_extension DESTINATION .)
+.. literalinclude:: ../../examples/python_packaging/CMakeLists.txt
+ :language: cmake
+ :start-after: [example.cmake.begin]
+ :end-before: [example.cmake.end]
Function ``tvm_ffi_configure_target`` sets up TVM-FFI include paths, link
against TVM-FFI library,
-generates stubs under the specified directory, and optionally debug symbols.
+generates stubs under ``STUB_DIR``, and can scaffold stub files when
``STUB_INIT`` is
+enabled.
Function ``tvm_ffi_install`` places necessary information, e.g. debug symbols
in macOS, next to
the shared library for proper packaging.
@@ -261,19 +244,10 @@ Define a :pep:`517` build backend in ``pyproject.toml``,
with the following step
- Specify the source directory of the package via ``wheel.packages``, and the
installation
destination via ``wheel.install-dir``.
-.. code-block:: toml
-
- [build-system]
- requires = ["scikit-build-core>=0.10.0", "apache-tvm-ffi"]
- build-backend = "scikit_build_core.build"
-
- [tool.scikit-build]
- # The wheel is Python ABI-agnostic
- wheel.py-api = "py3"
- # The package contains the Python module at `python/my_ffi_extension`
- wheel.packages = ["python/my_ffi_extension"]
- # The install dir matches the import name
- wheel.install-dir = "my_ffi_extension"
+.. literalinclude:: ../../examples/python_packaging/pyproject.toml
+ :language: toml
+ :start-after: [pyproject.build.begin]
+ :end-before: [pyproject.build.end]
Once fully specified, scikit-build-core will invoke CMake and drive the
extension building process.
diff --git a/examples/packaging/src/extension.cc
b/examples/packaging/src/extension.cc
deleted file mode 100644
index 38525c9..0000000
--- a/examples/packaging/src/extension.cc
+++ /dev/null
@@ -1,131 +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.
- */
-/*!
- * \file example.cc
- * \brief Example of a tvm-ffi based library that registers various functions.
- *
- * It is a simple example that demonstrates how to package a tvm-ffi library
into a python wheel.
- * The library is written in C++ and can be compiled into a shared library.
- * The shared library can then be loaded into python and used to call the
functions.
- */
-#include <tvm/ffi/container/tensor.h>
-#include <tvm/ffi/dtype.h>
-#include <tvm/ffi/error.h>
-#include <tvm/ffi/function.h>
-#include <tvm/ffi/reflection/registry.h>
-
-namespace my_ffi_extension {
-
-namespace ffi = tvm::ffi;
-
-/*!
- * \brief Raises a runtime error
- *
- * This is an example function to show how to raise and propagate
- * an error across the language boundary.
- *
- * \param msg The message to raise the error with
- */
-void RaiseError(ffi::String msg) { TVM_FFI_THROW(RuntimeError) << msg; }
-
-void AddOne(ffi::TensorView x, ffi::TensorView y) {
- // implementation of a library function
- TVM_FFI_ICHECK(x.ndim() == 1) << "x must be a 1D tensor";
- DLDataType f32_dtype{kDLFloat, 32, 1};
- TVM_FFI_ICHECK(x.dtype() == f32_dtype) << "x must be a float tensor";
- TVM_FFI_ICHECK(y.ndim() == 1) << "y must be a 1D tensor";
- TVM_FFI_ICHECK(y.dtype() == f32_dtype) << "y must be a float tensor";
- TVM_FFI_ICHECK(x.size(0) == y.size(0)) << "x and y must have the same shape";
- for (int i = 0; i < x.size(0); ++i) {
- static_cast<float*>(y.data_ptr())[i] =
static_cast<float*>(x.data_ptr())[i] + 1;
- }
-}
-
-// expose global symbol add_one
-TVM_FFI_DLL_EXPORT_TYPED_FUNC(add_one, my_ffi_extension::AddOne);
-
-/*!
- * \brief Example of a custom object that is exposed to the FFI library
- */
-class IntPairObj : public tvm::ffi::Object {
- public:
- int64_t a;
- int64_t b;
-
- IntPairObj() = default;
- IntPairObj(int64_t a, int64_t b) : a(a), b(b) {}
-
- int64_t GetFirst() const { return this->a; }
-
- // Required: declare type information
- TVM_FFI_DECLARE_OBJECT_INFO_FINAL("my_ffi_extension.IntPair", IntPairObj,
tvm::ffi::Object);
-};
-
-/*!
- * \brief Defines an explicit reference to IntPairObj
- *
- * A reference wrapper serves as a reference-counted ptr to the object.
- * you can use obj->field to access the fields of the object.
- */
-class IntPair : public tvm::ffi::ObjectRef {
- public:
- // Constructor
- explicit IntPair(int64_t a, int64_t b) { data_ =
tvm::ffi::make_object<IntPairObj>(a, b); }
-
- // Required: define object reference methods
- TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(IntPair, tvm::ffi::ObjectRef,
IntPairObj);
-};
-
-// The static initialization block is
-// called once when the library is loaded.
-TVM_FFI_STATIC_INIT_BLOCK() {
- namespace refl = tvm::ffi::reflection;
- // In this particular example, we use the reflection mechanisms to
- // register the functions directly into the global function table.
- //
- // This is an alternative approach to TVM_FFI_DLL_EXPORT_TYPED_FUNC
- // that exports the function directly as C symbol that follows tvm-ffi abi.
- //
- // - For functions that are expected to be static part of tvm_ffi_example
project,
- // one can use reflection mechanisms to register the globa function.
- // - For functions that are compiled and dynamically loaded at runtime,
consider
- // using the normal export mechanism so they won't be exposed to the
global function table.
- //
- // Make sure to have a unique name across all registered functions,
- // always prefix with a package namespace name to avoid name collision.
- //
- // The function can then be found via tvm_ffi.get_global_func(name)
- // If the function is expected to stay throughout the lifetime of the
program/
- //
- // When registering via reflection mechanisms, the library do not need to be
loaded via
- // tvm::ffi::Module::LoadFromFile, instead, just load the dll or simply
bundle into the
- // final project
- refl::GlobalDef().def("my_ffi_extension.raise_error", RaiseError);
- // register the object into the system
- // register field accessors and a global static function `__ffi_init__` as
ffi::Function
- refl::ObjectDef<IntPairObj>()
- .def(refl::init<int64_t, int64_t>())
- // Example static method that returns the second element of the pair
- .def_static("static_get_second", [](IntPair pair) -> int64_t { return
pair->b; })
- // Example to bind an instance method
- .def("get_first", &IntPairObj::GetFirst)
- .def_ro("a", &IntPairObj::a)
- .def_ro("b", &IntPairObj::b);
-}
-} // namespace my_ffi_extension
diff --git a/examples/packaging/CMakeLists.txt
b/examples/python_packaging/CMakeLists.txt
similarity index 92%
rename from examples/packaging/CMakeLists.txt
rename to examples/python_packaging/CMakeLists.txt
index 9b9fa80..d1f2879 100644
--- a/examples/packaging/CMakeLists.txt
+++ b/examples/python_packaging/CMakeLists.txt
@@ -23,7 +23,9 @@ find_package(
REQUIRED
)
find_package(tvm_ffi CONFIG REQUIRED)
+# [example.cmake.begin]
add_library(my_ffi_extension SHARED src/extension.cc)
tvm_ffi_configure_target(my_ffi_extension STUB_DIR "./python" STUB_INIT ON)
install(TARGETS my_ffi_extension DESTINATION .)
-tvm_ffi_install(my_ffi_extension)
+tvm_ffi_install(my_ffi_extension DESTINATION .)
+# [example.cmake.end]
diff --git a/examples/packaging/README.md b/examples/python_packaging/README.md
similarity index 74%
rename from examples/packaging/README.md
rename to examples/python_packaging/README.md
index 96d0939..e760710 100644
--- a/examples/packaging/README.md
+++ b/examples/python_packaging/README.md
@@ -31,7 +31,7 @@ packaging as well.
Use `uv pip` (the same tooling used in CI) to build and install the example
wheel:
```bash
-cd examples/packaging
+cd examples/python_packaging
uv pip install --reinstall --verbose .
```
@@ -42,20 +42,12 @@ Note: When running the auditwheel process, make sure to skip
## Run the example
-After installing the `my_ffi_extension` example package, you can run the
following example
-that invokes the `add_one` function exposed.
+After installing the `my_ffi_extension` example package, you can run the
following example.
```bash
python run_example.py
```
-This runs three flows: calling `add_one`, demonstrating `raise_error` with a
propagated traceback, and constructing/using the `IntPair` object.
-
-When possible, tvm_ffi will try to preserve backtrace across language
boundary. You will see output like
-
-```text
-File "src/extension.cc", line 45, in void
my_ffi_extension::RaiseError(tvm::ffi::String)
-```
-
-If you are in an IDE like VSCode, you can click and jump to the C++ lines of
error when
-the debug symbols are preserved.
+This runs four flows: calling `add_two` via the TVM-FFI ABI, calling `add_one`
via the global
+registry, calling `raise_error` to demonstrate error propagation, and
constructing/using the
+`IntPair` object.
diff --git a/examples/packaging/pyproject.toml
b/examples/python_packaging/pyproject.toml
similarity index 89%
rename from examples/packaging/pyproject.toml
rename to examples/python_packaging/pyproject.toml
index 7385120..7652bb4 100644
--- a/examples/packaging/pyproject.toml
+++ b/examples/python_packaging/pyproject.toml
@@ -33,14 +33,19 @@ requires-python = ">=3.9"
dependencies = ["apache-tvm-ffi"]
+# [pyproject.build.begin]
[build-system]
requires = ["scikit-build-core>=0.10.0", "apache-tvm-ffi"]
build-backend = "scikit_build_core.build"
[tool.scikit-build]
-# the wheel is abi agnostic
+# The wheel is Python ABI-agnostic
wheel.py-api = "py3"
-minimum-version = "build-system.requires"
+# The package contains the Python module at `python/my_ffi_extension`
+wheel.packages = ["python/my_ffi_extension"]
+# The install dir matches the import name
+wheel.install-dir = "my_ffi_extension"
+# [pyproject.build.end]
# Build configuration
build-dir = "build"
@@ -52,7 +57,4 @@ cmake.build-type = "RelWithDebInfo"
# Logging
logging.level = "INFO"
-
-# Wheel configuration
-wheel.packages = ["python/my_ffi_extension"]
-wheel.install-dir = "my_ffi_extension"
+minimum-version = "build-system.requires"
diff --git a/examples/packaging/python/my_ffi_extension/__init__.py
b/examples/python_packaging/python/my_ffi_extension/__init__.py
similarity index 100%
rename from examples/packaging/python/my_ffi_extension/__init__.py
rename to examples/python_packaging/python/my_ffi_extension/__init__.py
diff --git a/examples/packaging/python/my_ffi_extension/_ffi_api.py
b/examples/python_packaging/python/my_ffi_extension/_ffi_api.py
similarity index 94%
rename from examples/packaging/python/my_ffi_extension/_ffi_api.py
rename to examples/python_packaging/python/my_ffi_extension/_ffi_api.py
index be25187..0ea71b8 100644
--- a/examples/packaging/python/my_ffi_extension/_ffi_api.py
+++ b/examples/python_packaging/python/my_ffi_extension/_ffi_api.py
@@ -33,6 +33,7 @@ LIB = _FFI_LOAD_LIB("my_ffi_extension", "my_ffi_extension")
# fmt: off
_FFI_INIT_FUNC("my_ffi_extension", __name__)
if TYPE_CHECKING:
+ def add_one(_0: int, /) -> int: ...
def raise_error(_0: str, /) -> None: ...
# fmt: on
# tvm-ffi-stubgen(end)
@@ -49,9 +50,7 @@ class IntPair(_ffi_Object):
if TYPE_CHECKING:
@staticmethod
def __c_ffi_init__(_0: int, _1: int, /) -> Object: ...
- @staticmethod
- def static_get_second(_0: IntPair, /) -> int: ...
- def get_first(self, /) -> int: ...
+ def sum(self, /) -> int: ...
# fmt: on
# tvm-ffi-stubgen(end)
@@ -60,6 +59,7 @@ __all__ = [
# tvm-ffi-stubgen(begin): __all__
"LIB",
"IntPair",
+ "add_one",
"raise_error",
# tvm-ffi-stubgen(end)
]
diff --git a/examples/packaging/run_example.py
b/examples/python_packaging/run_example.py
similarity index 74%
rename from examples/packaging/run_example.py
rename to examples/python_packaging/run_example.py
index 94c90b0..fd3c03c 100644
--- a/examples/packaging/run_example.py
+++ b/examples/python_packaging/run_example.py
@@ -19,21 +19,23 @@
import traceback
import my_ffi_extension
-import torch
+
+
+def run_add_two() -> None:
+ """Invoke add_two from the extension and print the result."""
+ print("=========== Example 1: add_two ===========")
+ print(my_ffi_extension.LIB.add_two(1))
def run_add_one() -> None:
"""Invoke add_one from the extension and print the result."""
- print("=========== Example 1: add_one ===========")
- x = torch.tensor([1, 2, 3, 4, 5], dtype=torch.float32)
- y = torch.empty_like(x)
- my_ffi_extension.LIB.add_one(x, y)
- print(y)
+ print("=========== Example 2: add_one ===========")
+ print(my_ffi_extension.add_one(3))
def run_raise_error() -> None:
"""Invoke raise_error from the extension to demonstrate error handling."""
- print("=========== Example 2: raise_error ===========")
+ print("=========== Example 3: raise_error ===========")
try:
my_ffi_extension.raise_error("This is an error")
except RuntimeError:
@@ -42,13 +44,15 @@ def run_raise_error() -> None:
def run_int_pair() -> None:
"""Invoke IntPair from the extension to demonstrate object handling."""
- print("=========== Example 3: IntPair ===========")
+ print("=========== Example 4: IntPair ===========")
pair = my_ffi_extension.IntPair(1, 2)
- print(f"first={pair.get_first()}")
- print(f"second={my_ffi_extension.IntPair.static_get_second(pair)}")
+ print(f"a={pair.a}")
+ print(f"b={pair.b}")
+ print(f"sum={pair.sum()}")
if __name__ == "__main__":
+ run_add_two()
run_add_one()
run_raise_error()
run_int_pair()
diff --git a/examples/python_packaging/src/extension.cc
b/examples/python_packaging/src/extension.cc
new file mode 100644
index 0000000..ee66a04
--- /dev/null
+++ b/examples/python_packaging/src/extension.cc
@@ -0,0 +1,91 @@
+/*
+ * 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.
+ */
+/*!
+ * \file extension.cc
+ * \brief Example of a tvm-ffi based library that registers various functions.
+ */
+#include <tvm/ffi/function.h>
+#include <tvm/ffi/object.h>
+#include <tvm/ffi/reflection/registry.h>
+
+#include <cstdint>
+
+namespace my_ffi_extension {
+
+namespace ffi = tvm::ffi;
+
+// [tvm_ffi_abi.begin]
+static int AddTwo(int x) { return x + 2; }
+
+TVM_FFI_DLL_EXPORT_TYPED_FUNC(add_two, AddTwo);
+// [tvm_ffi_abi.end]
+
+// [global_function.begin]
+static int AddOne(int x) { return x + 1; }
+
+/*!
+ * \brief Raise a runtime error to demonstrate error propagation.
+ *
+ * \param msg The error message to raise.
+ *
+ * \code{.py}
+ * import my_ffi_extension
+ * try:
+ * my_ffi_extension.raise_error("boom")
+ * except RuntimeError:
+ * pass
+ * \endcode
+ */
+static void RaiseError(const ffi::String& msg) { TVM_FFI_THROW(RuntimeError)
<< msg; }
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
+ refl::GlobalDef()
+ .def("my_ffi_extension.add_one", AddOne)
+ .def("my_ffi_extension.raise_error", RaiseError);
+}
+// [global_function.end]
+
+// [object.begin]
+class IntPairObj : public ffi::Object {
+ public:
+ int64_t a;
+ int64_t b;
+
+ IntPairObj(int64_t a, int64_t b) : a(a), b(b) {}
+
+ int64_t Sum() const { return a + b; }
+
+ static constexpr bool _type_mutable = true;
+ TVM_FFI_DECLARE_OBJECT_INFO_FINAL(
+ /*type_key=*/"my_ffi_extension.IntPair",
+ /*class=*/IntPairObj,
+ /*parent_class=*/ffi::Object);
+};
+
+TVM_FFI_STATIC_INIT_BLOCK() {
+ namespace refl = tvm::ffi::reflection;
+ refl::ObjectDef<IntPairObj>()
+ .def(refl::init<int64_t, int64_t>())
+ .def_rw("a", &IntPairObj::a, "the first field")
+ .def_rw("b", &IntPairObj::b, "the second field")
+ .def("sum", &IntPairObj::Sum, "IntPairObj::Sum() method");
+}
+// [object.end]
+} // namespace my_ffi_extension
diff --git a/python/tvm_ffi/stub/cli.py b/python/tvm_ffi/stub/cli.py
index f2d34a9..1540cbc 100644
--- a/python/tvm_ffi/stub/cli.py
+++ b/python/tvm_ffi/stub/cli.py
@@ -263,163 +263,14 @@ def _parse_args() -> Options:
parser = argparse.ArgumentParser(
prog="tvm-ffi-stubgen",
description=(
- "Generate type stubs for TVM FFI extensions.\n\n"
- "In `--init-*` mode, it generates missing `_ffi_api.py` and
`__init__.py` files, "
- "based on the registered global functions and object types in the
loaded libraries.\n\n"
- "In normal mode, it processes the given files/directories
in-place, generating "
- "type stubs inside special `tvm-ffi-stubgen` blocks. Scroo down
for more details."
+ "Generate type stubs for TVM FFI extensions. It supports two
modes\n"
+ "- In `--init-*` mode, it generates missing `_ffi_api.py` and
`__init__.py` files, "
+ "based on the registered global functions and object types in the
loaded libraries.\n"
+ "- In normal mode, it processes the given files/directories
in-place, generating "
+ "type stubs inside special `tvm-ffi-stubgen` directive blocks.\n\n"
+ f"Documentation: {C.TERM_CYAN}{C.DOC_URL}{C.TERM_RESET}."
),
formatter_class=HelpFormatter,
- epilog=(
- "========\n"
- "Examples\n"
- "========\n\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 extension libraries\n"
- " tvm-ffi-stubgen --dlls
build/libmy_ext.so;build/libmy_2nd_ext.so my_pkg/_ffi_api.py\n\n"
- " # Package-level init (my-ffi-extension)\n"
- " tvm-ffi-stubgen examples/packaging/python \\\n"
- " --dlls examples/packaging/build/libmy_ffi_extension.dylib
\\\n"
- " --init-pypkg my-ffi-extension \\\n"
- " --init-lib my_ffi_extension \\\n"
- ' --init-prefix "my_ffi_extension."\n\n'
- "=====================\n"
- "Syntax of stub blocks\n"
- "=====================\n\n"
- "Global functions\n"
- "~~~~~~~~~~~~~~~~\n\n"
- " ```\n"
- f" {C.STUB_BEGIN} global/<registry-prefix>@<import-from
(default: tvm_ffi)>\n"
- f" {C.STUB_END}\n"
- " ```\n\n"
- "Generates TYPE_CHECKING-only stubs for functions in the global
registry under the prefix.\n\n"
- "Example:\n\n"
- " ```\n"
- f" {C.STUB_BEGIN} global/[email protected]\n"
- " # fmt: off\n"
- ' _FFI_INIT_FUNC("ffi", __name__)\n'
- " if TYPE_CHECKING:\n"
- " def Array(*args: Any) -> Any: ...\n"
- " def ArrayGetItem(_0: Sequence[Any], _1: int, /) -> Any:
...\n"
- " def ArraySize(_0: Sequence[Any], /) -> int: ...\n"
- " def Bytes(_0: bytes, /) -> bytes: ...\n"
- " ...\n"
- " def StructuralHash(_0: Any, _1: bool, _2: bool, /) ->
int: ...\n"
- " def SystemLib(*args: Any) -> Any: ...\n"
- " def ToJSONGraph(_0: Any, _1: Any, /) -> Any: ...\n"
- " def ToJSONGraphString(_0: Any, _1: Any, /) -> str: ...\n"
- " # fmt: on\n"
- f" {C.STUB_END}\n"
- " ```\n\n"
- "Objects\n"
- "~~~~~~~\n\n"
- " ```\n"
- f" {C.STUB_BEGIN} object/<type_key>\n"
- f" {C.STUB_END}\n"
- " ```\n\n"
- "Generates fields/methods for a class defined using TVM-FFI Object
APIs.\n\n"
- "Example:\n\n"
- " ```\n"
- ' @register_object("ffi.reflection.AccessPath")\n'
- " class AccessPath(tvm_ffi.Object):\n"
- f" {C.STUB_BEGIN} object/ffi.reflection.AccessPath\n"
- " # fmt: off\n"
- " parent: Object | None\n"
- " step: AccessStep | None\n"
- " depth: int\n"
- " if TYPE_CHECKING:\n"
- " @staticmethod\n"
- " def _root() -> AccessPath: ...\n"
- " def _extend(self, _1: AccessStep, /) -> AccessPath:
...\n"
- " def _attr(self, _1: str, /) -> AccessPath: ...\n"
- " def _array_item(self, _1: int, /) -> AccessPath:
...\n"
- " def _map_item(self, _1: Any, /) -> AccessPath: ...\n"
- " def _attr_missing(self, _1: str, /) -> AccessPath:
...\n"
- " def _array_item_missing(self, _1: int, /) ->
AccessPath: ...\n"
- " def _map_item_missing(self, _1: Any, /) ->
AccessPath: ...\n"
- " def _is_prefix_of(self, _1: AccessPath, /) -> bool:
...\n"
- " def _to_steps(self, /) -> Sequence[AccessStep]: ...\n"
- " def _path_equal(self, _1: AccessPath, /) -> bool:
...\n"
- " # fmt: on\n"
- f" {C.STUB_END}\n"
- " ```\n\n"
- "Import section\n"
- "~~~~~~~~~~~~~~\n\n"
- " ```\n"
- f" {C.STUB_BEGIN} import-section\n"
- " # fmt: off\n"
- " # isort: off\n"
- " from __future__ import annotations\n"
- " from ..registry import init_ffi_api as _FFI_INIT_FUNC\n"
- " from typing import TYPE_CHECKING\n"
- " if TYPE_CHECKING:\n"
- " from collections.abc import Mapping, Sequence\n"
- " from tvm_ffi import Device, Object, Tensor, dtype\n"
- " from tvm_ffi.testing import TestIntPair\n"
- " from typing import Any, Callable\n"
- " # isort: on\n"
- " # fmt: on\n"
- f" {C.STUB_END}\n"
- " ```\n\n"
- "Auto-populates imports used by generated stubs.\n\n"
- "Export\n"
- "~~~~~~\n\n"
- " ```\n"
- f" {C.STUB_BEGIN} export/_ffi_api\n"
- " # fmt: off\n"
- " # isort: off\n"
- " from ._ffi_api import * # noqa: F403\n"
- " from ._ffi_api import __all__ as _ffi_api__all__\n"
- ' if "__all__" not in globals():\n'
- " __all__ = []\n"
- " __all__.extend(_ffi_api__all__)\n"
- " # isort: on\n"
- " # fmt: on\n"
- f" {C.STUB_END}\n"
- " ```\n\n"
- "Re-exports a generated submodule's __all__ into the parent.\n\n"
- "__all__\n"
- "~~~~~~~\n\n"
- " ```\n"
- " __all__ = [\n"
- f" {C.STUB_BEGIN} __all__\n"
- ' "LIB",\n'
- ' "IntPair",\n'
- ' "raise_error",\n'
- f" {C.STUB_END}\n"
- " ]\n"
- " ```\n\n"
- "Populates __all__ with generated classes/functions and LIB (if
present).\n\n"
- "Type map\n"
- "~~~~~~~~\n\n"
- " ```\n"
- f" {C.STUB_TY_MAP} <type_key> -> <python_type>\n"
- " ```\n\n"
- "Maps runtime type keys to Python types used in generation.\n\n"
- "Example:\n\n"
- " ```\n"
- f" {C.STUB_TY_MAP} ffi.reflection.AccessStep ->
ffi.access_path.AccessStep\n"
- " ```\n\n"
- "Import object\n"
- "~~~~~~~~~~~~~\n\n"
- " ```\n"
- f" {C.STUB_IMPORT_OBJECT} <from>; <type_checking_only>;
<alias>\n"
- " ```\n\n"
- "Injects a custom import into generated code, optionally
TYPE_CHECKING-only.\n\n"
- "Example:\n\n"
- " ```\n"
- f" {C.STUB_IMPORT_OBJECT} ffi.Object;False;_ffi_Object\n"
- " ```\n\n"
- "Skip file\n"
- "~~~~~~~~~\n\n"
- " ```\n"
- f" {C.STUB_SKIP_FILE}\n"
- " ```\n\n"
- "Prevents stubgen from modifying the file."
- ),
)
parser.add_argument(
"--imports",
@@ -486,7 +337,7 @@ def _parse_args() -> Options:
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 "
+ "only .py and .pyi files are modified. Use tvm-ffi-stubgen
directives to "
"select where stubs are generated."
),
)
diff --git a/python/tvm_ffi/stub/consts.py b/python/tvm_ffi/stub/consts.py
index b6458e1..2beb388 100644
--- a/python/tvm_ffi/stub/consts.py
+++ b/python/tvm_ffi/stub/consts.py
@@ -47,6 +47,7 @@ TERM_BLUE = "\033[34m"
TERM_MAGENTA = "\033[35m"
TERM_CYAN = "\033[36m"
TERM_WHITE = "\033[37m"
+DOC_URL =
"https://tvm.apache.org/ffi/packaging/python_packaging.html#stub-generation-tool";
DEFAULT_SOURCE_EXTS = {".py", ".pyi"}
TY_MAP_DEFAULTS = {
