diff --git a/.github/workflows/ci_mainline_only.yml b/.github/workflows/ci_mainline_only.yml index 1912acbc..dbf170b9 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 c768768d..3166fdbf 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 `_. + +.. 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() - .def(refl::init()) - .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 38525c95..00000000 --- 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 -#include -#include -#include -#include - -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(y.data_ptr())[i] = static_cast(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(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() - .def(refl::init()) - // 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 9b9fa802..d1f28798 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 96d0939a..e7607100 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 7385120e..7652bb46 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 be25187b..0ea71b89 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 @@ # 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 @@ def get_first(self, /) -> int: ... # 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 94c90b0d..fd3c03cf 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 00000000..ee66a047 --- /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 +#include +#include + +#include + +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() + .def(refl::init()) + .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 f2d34a9f..1540cbc3 100644 --- a/python/tvm_ffi/stub/cli.py +++ b/python/tvm_ffi/stub/cli.py @@ -263,163 +263,14 @@ def _split_list_arg(arg: str | None) -> list[str]: 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/@\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/ffi@.registry\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/\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} -> \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} ; ; \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 _split_list_arg(arg: str | None) -> list[str]: 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 b6458e15..2beb388b 100644 --- a/python/tvm_ffi/stub/consts.py +++ b/python/tvm_ffi/stub/consts.py @@ -47,6 +47,7 @@ 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 = {