[Stubgen] Introduce --init-path for package-level generation

This PR introduces an additional flag `--init-path` to the tool
`tvm-ffi-stubgen`.

When this flag is specified, the tool will actively look for global
functions and objects that are not yet exported into Python, and
generate their type stubs into the path specified by `--init-path`.

To give an example, if a C++ TVM-FFI extension defines two items:
- A global function: `my_ffi_extension.add_one`
- An object: `my_ffi_extension.IntPair`
which are not yet exported to Python, the tool will generate the
following two files if `--init-path=/path/` is set:

```
/path/__init__.py
/path/_ffi_api.py
```

Where `__init__.py` contains:

```python
"""Package my_ffi_extension."""

/# tvm-ffi-stubgen(begin): export/_ffi_api
/# fmt: off
/# isort: off
from ._ffi_api import *  # noqa: F403
from ._ffi_api import __all__ as _ffi_api__all__
if "__all__" not in globals(): __all__ = []
__all__.extend(_ffi_api__all__)
/# isort: on
/# fmt: on
/# tvm-ffi-stubgen(end)
```

and `_ffi_api.py` contains:

```python
"""FFI API bindings for my_ffi_extension."""

/# tvm-ffi-stubgen(begin): import
/# fmt: off
/# isort: off
from __future__ import annotations
from typing import TYPE_CHECKING
if TYPE_CHECKING:
    from tvm_ffi import Object
/# isort: on
/# fmt: on
/# tvm-ffi-stubgen(end)

/# tvm-ffi-stubgen(begin): global/my_ffi_extension
/# fmt: off
/# isort: off
from tvm_ffi import init_ffi_api as _INIT
_INIT("my_ffi_extension", __name__)
/# isort: on
if TYPE_CHECKING:
    def raise_error(_0: str, /) -> None: ...
/# fmt: on
/# tvm-ffi-stubgen(end)

__all__ = [
    # tvm-ffi-stubgen(begin): __all__
    "IntPair",
    "raise_error",
    # tvm-ffi-stubgen(end)
]

/# isort: off
import tvm_ffi
/# isort: on

@tvm_ffi.register_object("my_ffi_extension.IntPair")
class IntPair(tvm_ffi.Object):
    """FFI binding for `my_ffi_extension.IntPair`."""

    /# tvm-ffi-stubgen(begin): object/my_ffi_extension.IntPair
    /# fmt: off
    a: int
    b: int
    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: ...
    /# fmt: on
    /# tvm-ffi-stubgen(end)
```

Author: junrushao

python/tvm_ffi/_ffi_api.py

@@ -29,10 +29,12 @@
 # fmt: on
 # tvm-ffi-stubgen(end)
 
-from . import registry
-
-# tvm-ffi-stubgen(begin): global/ffi
+# tvm-ffi-stubgen(begin): global/ffi@.registry
 # fmt: off
+# isort: off
+from .registry import init_ffi_api as _INIT
+_INIT("ffi", __name__)
+# isort: on
 if TYPE_CHECKING:
     def Array(*args: Any) -> Any: ...
     def ArrayGetItem(_0: Sequence[Any], _1: int, /) -> Any: ...
@@ -72,8 +74,6 @@ def ToJSONGraphString(_0: Any, _1: Any, /) -> str: ...
 # fmt: on
 # tvm-ffi-stubgen(end)
 
-registry.init_ffi_api("ffi", __name__)
-
 
 __all__ = [
     # tvm-ffi-stubgen(begin): __all__

python/tvm_ffi/_tensor.py

@@ -57,6 +57,11 @@ class Shape(tuple, PyNativeObject):
 
     _tvm_ffi_cached_object: Any
 
+    # tvm-ffi-stubgen(begin): object/ffi.Shape
+    # fmt: off
+    # fmt: on
+    # tvm-ffi-stubgen(end)
+
     def __new__(cls, content: tuple[int, ...]) -> Shape:
         if any(not isinstance(x, Integral) for x in content):
             raise ValueError("Shape must be a tuple of integers")

python/tvm_ffi/container.py

@@ -148,6 +148,11 @@ class Array(core.Object, Sequence[T]):
 
     """
 
+    # tvm-ffi-stubgen(begin): object/ffi.Array
+    # fmt: off
+    # fmt: on
+    # tvm-ffi-stubgen(end)
+
     def __init__(self, input_list: Iterable[T]) -> None:
         """Construct an Array from a Python sequence."""
         self.__init_handle_by_constructor__(_ffi_api.Array, *input_list)
@@ -290,6 +295,11 @@ class Map(core.Object, Mapping[K, V]):
 
     """
 
+    # tvm-ffi-stubgen(begin): object/ffi.Map
+    # fmt: off
+    # fmt: on
+    # tvm-ffi-stubgen(end)
+
     def __init__(self, input_dict: Mapping[K, V]) -> None:
         """Construct a Map from a Python mapping."""
         list_kvs: list[Any] = []

python/tvm_ffi/stub/analysis.py

@@ -18,6 +18,7 @@
 
 from __future__ import annotations
 
+from tvm_ffi._ffi_api import GetRegisteredTypeKeys
 from tvm_ffi.registry import list_global_func_names
 
 from . import consts as C
@@ -34,8 +35,27 @@ def collect_global_funcs() -> dict[str, list[FuncInfo]]:
         except ValueError:
             print(f"{C.TERM_YELLOW}[Skipped] Invalid name in global function: {name}{C.TERM_RESET}")
         else:
-            global_funcs.setdefault(prefix, []).append(FuncInfo.from_global_name(name))
+            try:
+                global_funcs.setdefault(prefix, []).append(FuncInfo.from_global_name(name))
+            except:
+                print(f"{C.TERM_YELLOW}[Skipped] Function has no type schema: {name}{C.TERM_RESET}")
     # Ensure stable ordering for deterministic output.
     for k in list(global_funcs.keys()):
         global_funcs[k].sort(key=lambda x: x.schema.name)
     return global_funcs
+
+
+def collect_type_keys() -> dict[str, list[str]]:
+    """Collect registered object type keys from TVM FFI's global registry."""
+    global_objects: dict[str, list[str]] = {}
+    for type_key in GetRegisteredTypeKeys():
+        try:
+            prefix, _ = type_key.rsplit(".", 1)
+        except ValueError:
+            pass
+        else:
+            global_objects.setdefault(prefix, []).append(type_key)
+    # Ensure stable ordering for deterministic output.
+    for k in list(global_objects.keys()):
+        global_objects[k].sort()
+    return global_objects

python/tvm_ffi/stub/cli.py

@@ -14,23 +14,23 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
-# tvm-ffi-stubgen(skip-file)
 """TVM-FFI Stub Generator (``tvm-ffi-stubgen``)."""
 
 from __future__ import annotations
 
 import argparse
 import ctypes
+import importlib
 import sys
 import traceback
 from pathlib import Path
 from typing import Callable
 
 from . import codegen as G
 from . import consts as C
-from .analysis import collect_global_funcs
+from .analysis import collect_global_funcs, collect_type_keys
 from .file_utils import FileInfo, collect_files
-from .utils import Options
+from .utils import FuncInfo, Options
 
 
 def _fn_ty_map(ty_map: dict[str, str], ty_used: set[str]) -> Callable[[str], str]:
@@ -55,71 +55,43 @@ def __main__() -> int:
     overview and examples of the block syntax.
     """
     opt = _parse_args()
+    for imp in opt.imports or []:
+        importlib.import_module(imp)
+    if opt.init_path:
+        opt.files.append(opt.init_path)
     dlls = [ctypes.CDLL(lib) for lib in opt.dlls]
     files: list[FileInfo] = collect_files([Path(f) for f in opt.files])
+    global_funcs: dict[str, list[FuncInfo]] = collect_global_funcs()
 
-    # Stage 1: Process `tvm-ffi-stubgen(ty-map)`
+    # Stage 1: Collect information
+    # - type maps: `tvm-ffi-stubgen(ty-map)`
+    # - defined global functions: `tvm-ffi-stubgen(begin): global/...`
+    # - defined object types: `tvm-ffi-stubgen(begin): object/...`
     ty_map: dict[str, str] = C.TY_MAP_DEFAULTS.copy()
-
-    def _stage_1(file: FileInfo) -> None:
-        for code in file.code_blocks:
-            if code.kind == "ty-map":
-                try:
-                    lhs, rhs = code.param.split("->")
-                except ValueError as e:
-                    raise ValueError(
-                        f"Invalid ty_map format at line {code.lineno_start}. Example: `A.B -> C.D`"
-                    ) from e
-                ty_map[lhs.strip()] = rhs.strip()
-
     for file in files:
         try:
-            _stage_1(file)
+            _stage_1(file, ty_map)
         except Exception:
             print(
                 f'{C.TERM_RED}[Failed] File "{file.path}": {traceback.format_exc()}{C.TERM_RESET}'
             )
 
-    # Stage 2: Process
+    # Stage 2. Generate stubs if they are not defined on the file.
+    if opt.init_path:
+        _stage_2(
+            files,
+            init_path=Path(opt.init_path).resolve(),
+            global_funcs=global_funcs,
+        )
+
+    # Stage 3: Process
     # - `tvm-ffi-stubgen(begin): global/...`
     # - `tvm-ffi-stubgen(begin): object/...`
-    global_funcs = collect_global_funcs()
-
-    def _stage_2(file: FileInfo) -> None:
-        all_defined = set()
+    for file in files:
         if opt.verbose:
             print(f"{C.TERM_CYAN}[File] {file.path}{C.TERM_RESET}")
-        ty_used: set[str] = set()
-        ty_on_file: set[str] = set()
-        fn_ty_map_fn = _fn_ty_map(ty_map, ty_used)
-        # Stage 2.1. Process `tvm-ffi-stubgen(begin): global/...`
-        for code in file.code_blocks:
-            if code.kind == "global":
-                funcs = global_funcs.get(code.param, [])
-                for func in funcs:
-                    all_defined.add(func.schema.name)
-                G.generate_global_funcs(code, funcs, fn_ty_map_fn, opt)
-        # Stage 2.2. Process `tvm-ffi-stubgen(begin): object/...`
-        for code in file.code_blocks:
-            if code.kind == "object":
-                type_key = code.param
-                ty_on_file.add(ty_map.get(type_key, type_key))
-                G.generate_object(code, fn_ty_map_fn, opt)
-        # Stage 2.3. Add imports for used types.
-        for code in file.code_blocks:
-            if code.kind == "import":
-                G.generate_imports(code, ty_used - ty_on_file, opt)
-                break  # Only one import block per file is supported for now.
-        # Stage 2.4. Add `__all__` for defined classes and functions.
-        for code in file.code_blocks:
-            if code.kind == "__all__":
-                G.generate_all(code, all_defined | ty_on_file, opt)
-                break  # Only one __all__ block per file is supported for now.
-        file.update(show_diff=opt.verbose, dry_run=opt.dry_run)
-
-    for file in files:
         try:
-            _stage_2(file)
+            _stage_3(file, opt, ty_map, global_funcs)
         except:
             print(
                 f'{C.TERM_RED}[Failed] File "{file.path}": {traceback.format_exc()}{C.TERM_RESET}'
@@ -128,6 +100,122 @@ def _stage_2(file: FileInfo) -> None:
     return 0
 
 
+def _stage_1(
+    file: FileInfo,
+    ty_map: dict[str, str],
+) -> None:
+    for code in file.code_blocks:
+        if code.kind == "ty-map":
+            try:
+                assert isinstance(code.param, str)
+                lhs, rhs = code.param.split("->")
+            except ValueError as e:
+                raise ValueError(
+                    f"Invalid ty_map format at line {code.lineno_start}. Example: `A.B -> C.D`"
+                ) from e
+            ty_map[lhs.strip()] = rhs.strip()
+
+
+def _stage_2(
+    files: list[FileInfo],
+    init_path: Path,
+    global_funcs: dict[str, list[FuncInfo]],
+) -> None:
+    def _find_or_insert_file(path: Path) -> FileInfo:
+        ret: FileInfo | None
+        if not path.exists():
+            ret = FileInfo(path=path, lines=(), code_blocks=[])
+        else:
+            for file in files:
+                if path.samefile(file.path):
+                    return file
+            ret = FileInfo.from_file(file=path)
+            assert ret is not None, f"Failed to read file: {path}"
+        files.append(ret)
+        return ret
+
+    # Step 0. Find out functions and classes already defined on files.
+    defined_func_prefixes: set[str] = {  # type: ignore[union-attr]
+        code.param[0] for file in files for code in file.code_blocks if code.kind == "global"
+    }
+    defined_objs: set[str] = {  # type: ignore[assignment]
+        code.param for file in files for code in file.code_blocks if code.kind == "object"
+    } | C.BUILTIN_TYPE_KEYS
+
+    # Step 1. Generate missing `_ffi_api.py` and `__init__.py` under each prefix.
+    prefixes: dict[str, list[str]] = collect_type_keys()
+    for prefix in global_funcs:
+        prefixes.setdefault(prefix, [])
+
+    for prefix, obj_names in prefixes.items():
+        if prefix.startswith("testing") or prefix.startswith("ffi"):
+            continue
+        funcs = sorted(
+            [] if prefix in defined_func_prefixes else global_funcs.get(prefix, []),
+            key=lambda f: f.schema.name,
+        )
+        objs = sorted(set(obj_names) - defined_objs)
+        if not funcs and not objs:
+            continue
+        # Step 1.1. Create target directory if not exists
+        directory = init_path / prefix.replace(".", "/")
+        directory.mkdir(parents=True, exist_ok=True)
+        # Step 1.2. Generate `_ffi_api.py`
+        target_path = directory / "_ffi_api.py"
+        target_file = _find_or_insert_file(target_path)
+        with target_path.open("a", encoding="utf-8") as f:
+            f.write(G.generate_ffi_api(target_file.code_blocks, prefix, objs))
+        target_file.reload()
+        # Step 1.3. Generate `__init__.py`
+        target_path = directory / "__init__.py"
+        target_file = _find_or_insert_file(target_path)
+        with target_path.open("a", encoding="utf-8") as f:
+            f.write(G.generate_init(target_file.code_blocks, prefix, submodule="_ffi_api"))
+        target_file.reload()
+
+
+def _stage_3(
+    file: FileInfo,
+    opt: Options,
+    ty_map: dict[str, str],
+    global_funcs: dict[str, list[FuncInfo]],
+) -> None:
+    all_defined = set()
+    ty_used: set[str] = set()
+    ty_on_file: set[str] = set()
+    fn_ty_map_fn = _fn_ty_map(ty_map, ty_used)
+    # Stage 2.1. Process `tvm-ffi-stubgen(begin): global/...`
+    for code in file.code_blocks:
+        if code.kind == "global":
+            funcs = global_funcs.get(code.param[0], [])
+            for func in funcs:
+                all_defined.add(func.schema.name)
+            G.generate_global_funcs(code, funcs, fn_ty_map_fn, opt)
+    # Stage 2.2. Process `tvm-ffi-stubgen(begin): object/...`
+    for code in file.code_blocks:
+        if code.kind == "object":
+            type_key = code.param
+            assert isinstance(type_key, str)
+            ty_on_file.add(ty_map.get(type_key, type_key))
+            G.generate_object(code, fn_ty_map_fn, opt)
+    # Stage 2.3. Add imports for used types.
+    for code in file.code_blocks:
+        if code.kind == "import":
+            G.generate_imports(code, ty_used - ty_on_file, opt)
+            break  # Only one import block per file is supported for now.
+    # Stage 2.4. Add `__all__` for defined classes and functions.
+    for code in file.code_blocks:
+        if code.kind == "__all__":
+            G.generate_all(code, all_defined | ty_on_file, opt)
+            break  # Only one __all__ block per file is supported for now.
+    # Stage 2.5. Process `tvm-ffi-stubgen(begin): export/...`
+    for code in file.code_blocks:
+        if code.kind == "export":
+            G.generate_export(code)
+    # Finalize: write back to file
+    file.update(verbose=opt.verbose, dry_run=opt.dry_run)
+
+
 def _parse_args() -> Options:
     class HelpFormatter(argparse.ArgumentDefaultsHelpFormatter, argparse.RawTextHelpFormatter):
         pass
@@ -149,16 +237,16 @@ class HelpFormatter(argparse.ArgumentDefaultsHelpFormatter, argparse.RawTextHelp
             "  # 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"
+            f"  {C.STUB_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"
+            f"  {C.STUB_END}\n\n"
+            f"  {C.STUB_BEGIN} object/<type_key>\n"
+            f"  {C.STUB_TY_MAP}: list -> Sequence\n"
+            f"  {C.STUB_TY_MAP}: dict -> Mapping\n"
             "  ... generated fields and methods ...\n"
-            "  # tvm-ffi-stubgen(end)\n\n"
+            f"  {C.STUB_END}\n\n"
             "  # Skip a file entirely\n"
-            "  # tvm-ffi-stubgen(skip-file)\n\n"
+            f"  {C.STUB_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"
@@ -167,6 +255,12 @@ class HelpFormatter(argparse.ArgumentDefaultsHelpFormatter, argparse.RawTextHelp
             "    is provided by native extensions.\n"
         ),
     )
+    parser.add_argument(
+        "--imports",
+        nargs="*",
+        metavar="IMPORTS",
+        help=("Additional imports to load before generation."),
+    )
     parser.add_argument(
         "--dlls",
         nargs="*",
@@ -179,13 +273,19 @@ class HelpFormatter(argparse.ArgumentDefaultsHelpFormatter, argparse.RawTextHelp
         ),
         default=[],
     )
+    parser.add_argument(
+        "--init-path",
+        type=str,
+        default="",
+        help="If specified, generate stubs under the given package prefix.",
+    )
     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."
+            f"indentation of the corresponding '{C.STUB_BEGIN}' line."
         ),
     )
     parser.add_argument(