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 17ff30b doc: Auto-extract type annotations from function definition
(#106)
17ff30b is described below
commit 17ff30bc77a9201e0a190431be8cf9376075e7a5
Author: Junru Shao <[email protected]>
AuthorDate: Mon Oct 13 04:37:33 2025 -0700
doc: Auto-extract type annotations from function definition (#106)
A new plugin `sphinx-autodoc-typehints` is introduced that helps
automatically extract type annotations for function definitions.
It makes it possible, according to DRY (don't repeat yourself) rule, to
avoid defining type annotations the 2nd time inside docstring, which is
very much prone to inconsistency.
---
docs/conf.py | 17 ++++++++++++++-
pyproject.toml | 1 +
python/tvm_ffi/_convert.py | 4 ++--
python/tvm_ffi/_dtype.py | 4 ++--
python/tvm_ffi/access_path.py | 16 +++++++-------
python/tvm_ffi/container.py | 20 +++++++++---------
python/tvm_ffi/cpp/load_inline.py | 18 ++++++++--------
python/tvm_ffi/dataclasses/c_class.py | 4 ++--
python/tvm_ffi/dataclasses/field.py | 6 +++---
python/tvm_ffi/error.py | 20 +++++++++---------
python/tvm_ffi/libinfo.py | 6 +++---
python/tvm_ffi/module.py | 40 +++++++++++++++++------------------
python/tvm_ffi/registry.py | 28 ++++++++++++------------
python/tvm_ffi/serialization.py | 10 ++++-----
python/tvm_ffi/stream.py | 18 ++++++++--------
python/tvm_ffi/testing.py | 6 +++---
16 files changed, 117 insertions(+), 101 deletions(-)
diff --git a/docs/conf.py b/docs/conf.py
index b380a84..cc421af 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -55,6 +55,7 @@ extensions = [
"sphinx.ext.napoleon",
"sphinx.ext.viewcode",
"sphinx.ext.ifconfig",
+ "sphinx_autodoc_typehints",
"sphinx_copybutton",
"sphinx_reredirects",
"sphinx_tabs.tabs",
@@ -178,8 +179,22 @@ autodoc_default_options = {
"inherited-members": False,
"member-order": "bysource",
}
-autodoc_typehints = "description"
+autodoc_typehints = "description" # or "none"
+always_use_bars_union = True
+# Preserve how defaults are written in your source (e.g., DEFAULT_SENTINEL)
+# Requires Sphinx ≥ 4.0
+autodoc_preserve_defaults = True
+
+# Ask the extension to include defaults alongside types
+# 'braces' works well with NumPy-style "Parameters" tables
+typehints_defaults = "comma" # also accepts: "comma", "braces-after"
+
+# Optional: also add stubs for params you didn't list in the docstring
+always_document_param_types = True
+
+# Optional (pairs nicely with NumPy style)
+napoleon_use_rtype = False
# -- Other Options --------------------------------------------------------
templates_path = []
diff --git a/pyproject.toml b/pyproject.toml
index 24444eb..a76862d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -64,6 +64,7 @@ docs = [
"sphinx-reredirects==0.1.2",
"sphinx-tabs==3.4.1",
"sphinx-toolbox==3.4.0",
+ "sphinx-autodoc-typehints",
"sphinxcontrib-mermaid",
"sphinxcontrib-napoleon==0.7",
"sphinxcontrib_httpdomain==1.8.1",
diff --git a/python/tvm_ffi/_convert.py b/python/tvm_ffi/_convert.py
index 05e99da..1ff91af 100644
--- a/python/tvm_ffi/_convert.py
+++ b/python/tvm_ffi/_convert.py
@@ -42,12 +42,12 @@ def convert(value: Any) -> Any: # noqa: PLR0911,PLR0912
Parameters
----------
- value : Any
+ value
The python object to be converted.
Returns
-------
- ffi_obj : Any
+ ffi_obj
The converted TVM FFI object.
Note
diff --git a/python/tvm_ffi/_dtype.py b/python/tvm_ffi/_dtype.py
index 8306593..3aae074 100644
--- a/python/tvm_ffi/_dtype.py
+++ b/python/tvm_ffi/_dtype.py
@@ -49,7 +49,7 @@ class dtype(str):
Parameters
----------
- dtype_str : str
+ dtype_str
Note
----
@@ -77,7 +77,7 @@ class dtype(str):
Parameters
----------
- lanes : int
+ lanes
The number of lanes.
Returns
diff --git a/python/tvm_ffi/access_path.py b/python/tvm_ffi/access_path.py
index c906347..a048a8a 100644
--- a/python/tvm_ffi/access_path.py
+++ b/python/tvm_ffi/access_path.py
@@ -106,7 +106,7 @@ class AccessPath(Object):
Parameters
----------
- other : AccessPath
+ other
The access path to check if it is a prefix of this access path
Returns
@@ -122,7 +122,7 @@ class AccessPath(Object):
Parameters
----------
- attr_key : str
+ attr_key
The key of the attribute to access
Returns
@@ -138,7 +138,7 @@ class AccessPath(Object):
Parameters
----------
- attr_key : str
+ attr_key
The key of the attribute to access
Returns
@@ -154,7 +154,7 @@ class AccessPath(Object):
Parameters
----------
- index : int
+ index
The index of the item to access
Returns
@@ -170,7 +170,7 @@ class AccessPath(Object):
Parameters
----------
- index : int
+ index
The index of the item to access
Returns
@@ -186,7 +186,7 @@ class AccessPath(Object):
Parameters
----------
- key : Any
+ key
The key of the item to access
Returns
@@ -202,7 +202,7 @@ class AccessPath(Object):
Parameters
----------
- key : Any
+ key
The key of the item to access
Returns
@@ -218,7 +218,7 @@ class AccessPath(Object):
Returns
-------
- List[AccessStep]
+ access_steps
The list of access steps
"""
diff --git a/python/tvm_ffi/container.py b/python/tvm_ffi/container.py
index d4aa24f..a1b51a4 100644
--- a/python/tvm_ffi/container.py
+++ b/python/tvm_ffi/container.py
@@ -48,21 +48,21 @@ def getitem_helper(
Parameters
----------
- obj: Any
+ obj
The original object
- elem_getter : Callable[[Any, int], T]
+ elem_getter
A simple function that takes index and return a single element.
- length : int
+ length
The size of the array
- idx : SupportsIndex or slice
+ idx
The argument passed to getitem
Returns
-------
- result : object
+ result
The element for integer indices or a ``list`` for slices.
"""
@@ -90,7 +90,7 @@ class Array(core.Object, Sequence[T]):
Parameters
----------
- input_list : Iterable[T]
+ input_list
The list of values to be stored in the array.
See Also
@@ -231,7 +231,7 @@ class Map(core.Object, Mapping[K, V]):
Parameters
----------
- input_dict : Mapping[K, V]
+ input_dict
The dictionary of values to be stored in the map.
See Also
@@ -299,15 +299,15 @@ class Map(core.Object, Mapping[K, V]):
Parameters
----------
- key : object
+ key
The attribute key.
- default : object
+ default
The default object.
Returns
-------
- value: object
+ value
The result value.
"""
diff --git a/python/tvm_ffi/cpp/load_inline.py
b/python/tvm_ffi/cpp/load_inline.py
index 623ff88..260ccc2 100644
--- a/python/tvm_ffi/cpp/load_inline.py
+++ b/python/tvm_ffi/cpp/load_inline.py
@@ -399,40 +399,40 @@ def build_inline( # noqa: PLR0915, PLR0912
Parameters
----------
- name: str
+ name
The name of the tvm ffi module.
- cpp_sources: Sequence[str] | str, optional
+ cpp_sources
The C++ source code. It can be a list of sources or a single source.
- cuda_sources: Sequence[str] | str, optional
+ cuda_sources
The CUDA source code. It can be a list of sources or a single source.
- functions: Mapping[str, str] | Sequence[str] | str, optional
+ functions
The functions in cpp_sources or cuda_source that will be exported to
the tvm ffi module. When a mapping is
given, the keys are the names of the exported functions, and the
values are docstrings for the functions. When
a sequence or a single string is given, they are the functions needed
to be exported, and the docstrings are set
to empty strings. A single function name can also be given as a
string. When cpp_sources is given, the functions
must be declared (not necessarily defined) in the cpp_sources. When
cpp_sources is not given, the functions
must be defined in the cuda_sources. If not specified, no function
will be exported.
- extra_cflags: Sequence[str], optional
+ extra_cflags
The extra compiler flags for C++ compilation.
The default flags are:
- On Linux/macOS: ['-std=c++17', '-fPIC', '-O2']
- On Windows: ['/std:c++17', '/O2']
- extra_cuda_cflags: Sequence[str], optional
+ extra_cuda_cflags
The extra compiler flags for CUDA compilation.
- extra_ldflags: Sequence[str], optional
+ extra_ldflags
The extra linker flags.
The default flags are:
- On Linux/macOS: ['-shared']
- On Windows: ['/DLL']
- extra_include_paths: Sequence[str], optional
+ extra_include_paths
The extra include paths.
- build_directory: str, optional
+ build_directory
The build directory. If not specified, a default tvm ffi cache
directory will be used. By default, the
cache directory is ``~/.cache/tvm-ffi``. You can also set the
``TVM_FFI_CACHE_DIR`` environment variable to
specify the cache directory.
diff --git a/python/tvm_ffi/dataclasses/c_class.py
b/python/tvm_ffi/dataclasses/c_class.py
index 42dc4fd..35c22d7 100644
--- a/python/tvm_ffi/dataclasses/c_class.py
+++ b/python/tvm_ffi/dataclasses/c_class.py
@@ -57,12 +57,12 @@ def c_class(
Parameters
----------
- type_key : str
+ type_key
The reflection key that identifies the C++ type in the FFI registry,
e.g. ``"testing.MyClass"`` as registered in
``src/ffi/extra/testing.cc``.
- init : bool, default True
+ init
If ``True`` and the Python class does not define ``__init__``, an
initializer is auto-generated that mirrors the reflected constructor
signature. The generated initializer calls the C++ ``__init__``
diff --git a/python/tvm_ffi/dataclasses/field.py
b/python/tvm_ffi/dataclasses/field.py
index 92a28cb..b1d7d78 100644
--- a/python/tvm_ffi/dataclasses/field.py
+++ b/python/tvm_ffi/dataclasses/field.py
@@ -67,15 +67,15 @@ def field(
Parameters
----------
- default : Any, optional
+ default
A literal default value that populates the field when no argument
is given. At most one of ``default`` or ``default_factory`` may be
given.
- default_factory : Callable[[], Any], optional
+ default_factory
A zero-argument callable that produces the default. This matches the
semantics of :func:`dataclasses.field` and is useful for mutable
defaults such as ``list`` or ``dict``.
- init : bool, default True
+ init
If ``True`` the field is included in the generated ``__init__``.
If ``False`` the field is omitted from input arguments of ``__init__``.
diff --git a/python/tvm_ffi/error.py b/python/tvm_ffi/error.py
index fd0bf2b..14dd2b0 100644
--- a/python/tvm_ffi/error.py
+++ b/python/tvm_ffi/error.py
@@ -33,12 +33,12 @@ def _parse_backtrace(backtrace: str) -> list[tuple[str,
int, str]]:
Parameters
----------
- backtrace : str
+ backtrace
The backtrace string.
Returns
-------
- result : List[Tuple[str, int, str]]
+ result
The list of (filename, lineno, func)
"""
@@ -106,18 +106,18 @@ class TracebackManager:
Parameters
----------
- tb : types.TracebackType
+ tb
The traceback to append to.
- filename : str
+ filename
The filename of the traceback
- lineno : int
+ lineno
The line number of the traceback
- func : str
+ func
The function name of the traceback
Returns
-------
- new_tb : types.TracebackType
+ new_tb
The new traceback with the appended frame.
"""
@@ -163,15 +163,15 @@ def register_error(
Parameters
----------
- name_or_cls : str or class
+ name_or_cls
The name of the error class.
- cls : class
+ cls
The class to register.
Returns
-------
- fregister : function
+ fregister
Register function if f is not specified.
Examples
diff --git a/python/tvm_ffi/libinfo.py b/python/tvm_ffi/libinfo.py
index 8d92df3..fd0e784 100644
--- a/python/tvm_ffi/libinfo.py
+++ b/python/tvm_ffi/libinfo.py
@@ -26,15 +26,15 @@ def split_env_var(env_var: str, split: str) -> list[str]:
Parameters
----------
- env_var : str
+ env_var
Name of environment variable.
- split : str
+ split
String to split env_var on.
Returns
-------
- splits : list(string)
+ splits
If env_var exists, split env_var. Otherwise, empty list.
"""
diff --git a/python/tvm_ffi/module.py b/python/tvm_ffi/module.py
index 9db31a1..ff05f93 100644
--- a/python/tvm_ffi/module.py
+++ b/python/tvm_ffi/module.py
@@ -76,7 +76,7 @@ class Module(core.Object):
Returns
-------
- modules : list of Module
+ modules
The module
"""
@@ -95,15 +95,15 @@ class Module(core.Object):
Parameters
----------
- name : str
+ name
The name of the function
- query_imports : bool
+ query_imports
Whether to also query modules imported by this module.
Returns
-------
- b : Bool
+ b
True if module (or one of its imports) has a definition for name.
"""
@@ -123,15 +123,15 @@ class Module(core.Object):
Parameters
----------
- name : str
+ name
The name of the function
- query_imports : bool
+ query_imports
Whether also query modules imported by this module.
Returns
-------
- f : tvm_ffi.Function
+ f
The result function.
"""
@@ -146,7 +146,7 @@ class Module(core.Object):
Parameters
----------
- module : tvm.runtime.Module
+ module
The other module.
"""
@@ -168,12 +168,12 @@ class Module(core.Object):
Parameters
----------
- fmt : str, optional
+ fmt
The specified format.
Returns
-------
- source : str
+ source
The result source code.
"""
@@ -188,7 +188,7 @@ class Module(core.Object):
Returns
-------
- mask : int
+ mask
Bitmask of runtime module property
"""
@@ -199,7 +199,7 @@ class Module(core.Object):
Returns
-------
- b : Bool
+ b
True if the module is binary serializable.
"""
@@ -210,7 +210,7 @@ class Module(core.Object):
Returns
-------
- b : Bool
+ b
True if the module is runnable.
"""
@@ -223,7 +223,7 @@ class Module(core.Object):
Returns
-------
- b : Bool
+ b
True if the module is compilation exportable.
"""
@@ -238,9 +238,9 @@ class Module(core.Object):
Parameters
----------
- file_name : str
+ file_name
The name of the file.
- fmt : str
+ fmt
The format of the file.
See Also
@@ -263,13 +263,13 @@ def system_lib(symbol_prefix: str = "") -> Module:
Parameters
----------
- symbol_prefix: str = ""
+ symbol_prefix
Optional symbol prefix that can be used for search. When we lookup a
symbol
symbol_prefix + name will first be searched, then the name without
symbol_prefix.
Returns
-------
- module : runtime.Module
+ module
The system-wide library module.
"""
@@ -281,12 +281,12 @@ def load_module(path: str) -> Module:
Parameters
----------
- path : str
+ path
The path to the module file.
Returns
-------
- module : :py:class:`tvm_ffi.Module`
+ module
The loaded module
Examples
diff --git a/python/tvm_ffi/registry.py b/python/tvm_ffi/registry.py
index 6a9dad6..3080cce 100644
--- a/python/tvm_ffi/registry.py
+++ b/python/tvm_ffi/registry.py
@@ -34,7 +34,7 @@ def register_object(type_key: str | type | None = None) ->
Callable[[type], type
Parameters
----------
- type_key : str or cls
+ type_key
The type key of the node
Examples
@@ -88,18 +88,18 @@ def register_global_func(
Parameters
----------
- func_name : str or function
+ func_name
The function name
- f : function, optional
+ f
The function to be registered.
- override: boolean optional
+ override
Whether override existing entry.
Returns
-------
- fregister : function
+ fregister
Register function if f is not specified.
Examples
@@ -156,15 +156,15 @@ def get_global_func(name: str, allow_missing: bool =
False) -> core.Function | N
Parameters
----------
- name : str
+ name
The name of the global function
- allow_missing : bool
+ allow_missing
Whether allow missing function or raise an error.
Returns
-------
- func : Function
+ func
The function to be returned, None if function is missing.
See Also
@@ -180,7 +180,7 @@ def list_global_func_names() -> list[str]:
Returns
-------
- names : list
+ names
List of global functions names.
"""
@@ -194,7 +194,7 @@ def remove_global_func(name: str) -> None:
Parameters
----------
- name : str
+ name
The name of the global function
"""
@@ -206,12 +206,12 @@ def get_global_func_metadata(name: str) -> dict[str, Any]:
Parameters
----------
- name : str
+ name
The name of the global function
Returns
-------
- metadata : dict
+ metadata
The metadata of the function
"""
@@ -223,10 +223,10 @@ def init_ffi_api(namespace: str, target_module_name: str
| None = None) -> None:
Parameters
----------
- namespace : str
+ namespace
The namespace of the source registry
- target_module_name : str
+ target_module_name
The target module name if different from namespace
Examples
diff --git a/python/tvm_ffi/serialization.py b/python/tvm_ffi/serialization.py
index 6960eda..285d764 100644
--- a/python/tvm_ffi/serialization.py
+++ b/python/tvm_ffi/serialization.py
@@ -32,15 +32,15 @@ def to_json_graph_str(obj: Any, metadata: dict[str, Any] |
None = None) -> str:
Parameters
----------
- obj : Any
+ obj
The object to save.
- metadata : Optional[dict], optional
+ metadata
Extra metadata to save into the json graph string.
Returns
-------
- json_str : str
+ json_str
The JSON graph string.
"""
@@ -55,12 +55,12 @@ def from_json_graph_str(json_str: str) -> Any:
Parameters
----------
- json_str : str
+ json_str
The JSON graph string to load.
Returns
-------
- obj : Any
+ obj
The loaded object.
"""
diff --git a/python/tvm_ffi/stream.py b/python/tvm_ffi/stream.py
index e00d422..b2083f3 100644
--- a/python/tvm_ffi/stream.py
+++ b/python/tvm_ffi/stream.py
@@ -34,10 +34,10 @@ class StreamContext:
Parameters
----------
- device : Device
+ device
The device to which the stream belongs.
- stream : Union[int, c_void_p]
+ stream
The stream handle.
See Also
@@ -100,12 +100,12 @@ try:
Parameters
----------
- context : Any = None
+ context
The wrapped torch stream or cuda graph.
Returns
-------
- context : tvm_ffi.TorchStreamContext
+ context
The ffi stream context wrapping torch stream context.
Examples
@@ -139,15 +139,15 @@ def use_raw_stream(device: core.Device, stream:
Union[int, c_void_p]) -> StreamC
Parameters
----------
- device : tvm_ffi.Device
+ device
The device to which the stream belongs.
- stream : Union[int, c_void_p]
+ stream
The stream handle.
Returns
-------
- context : tvm_ffi.StreamContext
+ context
The ffi stream context.
Note
@@ -168,12 +168,12 @@ def get_raw_stream(device: core.Device) -> int:
Parameters
----------
- device : tvm_ffi.Device
+ device
The device to which the stream belongs.
Returns
-------
- stream : int
+ stream
The current ffi stream.
"""
diff --git a/python/tvm_ffi/testing.py b/python/tvm_ffi/testing.py
index 2ea89ce..71467df 100644
--- a/python/tvm_ffi/testing.py
+++ b/python/tvm_ffi/testing.py
@@ -114,14 +114,14 @@ def create_object(type_key: str, **kwargs: Any) -> Object:
Parameters
----------
- type_key : str
+ type_key
The type key of the object.
- kwargs : dict
+ kwargs
The keyword arguments to the object.
Returns
-------
- obj : object
+ obj
The created object.
Note
