junrushao commented on code in PR #101: URL: https://github.com/apache/tvm-ffi/pull/101#discussion_r2424490404
########## python/tvm_ffi/stub/stubgen.py: ########## @@ -0,0 +1,517 @@ +# 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. +"""TVM-FFI Stub Generator (``tvm-ffi-stubgen``). + +Overview +-------- +This module powers the ``tvm-ffi-stubgen`` command line tool which generates +in-place type stubs for Python modules that integrate with the TVM FFI +ecosystem. It scans ``.py``/``.pyi`` files for special comment markers and +fills the enclosed blocks with precise, static type annotations derived from +runtime metadata exposed by TVM FFI. + +Why you might use this +---------------------- +- You author a Python module that binds to C++/C via TVM FFI and want + high-quality type hints for functions, objects, and methods. +- You maintain a downstream extension that registers global functions or + FFI object types and want your Python API surface to be type-checker + friendly without manually writing stubs. + +How it works (in one sentence) +------------------------------ +``tvm-ffi-stubgen`` replaces the content between special ``tvm-ffi-stubgen`` markers +with generated code guarded by ``if TYPE_CHECKING: ...`` so that the runtime +behavior is unchanged while static analyzers get rich types. + +Stub block markers +------------------ +Insert one of the following begin/end markers in your source, then run +``tvm-ffi-stubgen``. Indentation on the ``begin`` line is preserved; generated +content is additionally indented by ``--indent`` spaces (default: 4). + +1) Global function stubs + + Mark all global functions whose names start with a registry prefix + (e.g. ``ffi`` or ``my_ffi_extension``): + + .. code-block:: python + + from typing import TYPE_CHECKING + + # tvm-ffi-stubgen(begin): global/ffi + if TYPE_CHECKING: + # fmt: off + # (generated by tvm-ffi-stubgen) + # fmt: on + # tvm-ffi-stubgen(end) + + ``tvm-ffi-stubgen`` expands this with function signatures discovered via the + TVM FFI global function registry. + +2) Object type stubs + + Mark fields and methods for a registered FFI object type using its + ``type_key`` (the key passed to ``@register_object``): + + .. code-block:: python + + @register_object("testing.SchemaAllTypes") + class _SchemaAllTypes: + # tvm-ffi-stubgen(begin): object/testing.SchemaAllTypes + # tvm-ffi-stubgen(ty_map): testing.SchemaAllTypes -> _SchemaAllTypes + if TYPE_CHECKING: + # fmt: off + # (generated by tvm-ffi-stubgen) + # fmt: on + # tvm-ffi-stubgen(end) + + ``tvm-ffi-stubgen`` expands this with annotated attributes and method stub + signatures. The special C FFI initializer ``__ffi_init__`` is exposed as + ``__c_ffi_init__`` to avoid interfering with your Python ``__init__``. + +3) Skip whole file + + If a source file should never be modified by the stub generator, add the + following directive anywhere in the file: + + .. code-block:: python + + # tvm-ffi-stubgen(skip-file) + + When present, ``tvm-ffi-stubgen`` skips processing this file entirely. This + is useful for files that are generated by other tooling or vendored. + +Optional type mapping lines +--------------------------- +Inside a stub block you may add mapping hints to rename fully-qualified type +names to simpler aliases in the generated output: + +.. code-block:: python + + # tvm-ffi-stubgen(ty_map): A.B.C -> C + # tvm-ffi-stubgen(ty_map): list -> Sequence + # tvm-ffi-stubgen(ty_map): dict -> Mapping + +By default, ``list`` is shown as ``Sequence`` and ``dict`` as ``Mapping``. +If you use names such as ``Sequence``/``Mapping``, ensure they are available +to type checkers in your module, for example: + +.. code-block:: python + + from typing import TYPE_CHECKING + if TYPE_CHECKING: + from collections.abc import Mapping, Sequence + +Runtime requirements +-------------------- +- Python must be able to import ``tvm_ffi``. +- The process needs access to the TVM runtime and any extension libraries that + provide the global functions or object types you want to stub. Use the + ``--dlls`` option to preload shared libraries when necessary. + +What files are modified +----------------------- +Only files with extensions ``.py`` and ``.pyi`` are scanned. Files are updated +in place. A colored unified diff is printed for each change. + +CLI quick start +--------------- + +.. code-block:: bash + + # Generate stubs for a single file + tvm-ffi-stubgen python/tvm_ffi/_ffi_api.py + + # Recursively scan directories for tvm-ffi-stubgen blocks + tvm-ffi-stubgen python/tvm_ffi examples/packaging/python/my_ffi_extension + + # Preload TVM runtime and your extension library before generation + tvm-ffi-stubgen \ + --dlls build/libtvm_runtime.dylib build/libmy_ext.dylib \ + python/tvm_ffi/_ffi_api.py + +Exit status +----------- +Returns 0 on success and 1 if any file fails to process. + +""" + +from __future__ import annotations + +import argparse +import ctypes +import dataclasses +import difflib +import logging +import sys +from io import StringIO +from pathlib import Path +from typing import Callable + +from tvm_ffi.core import TypeSchema, _lookup_or_register_type_info_from_type_key +from tvm_ffi.registry import get_global_func_metadata, list_global_func_names + +DEFAULT_SOURCE_EXTS = {".py", ".pyi"} +STUB_BEGIN = "# tvm-ffi-stubgen(begin):" +STUB_END = "# tvm-ffi-stubgen(end)" +STUB_TY_MAP = "# tvm-ffi-stubgen(ty_map):" +STUB_SKIP_FILE = "# tvm-ffi-stubgen(skip-file)" + +TERM_RESET = "\033[0m" +TERM_BOLD = "\033[1m" +TERM_RED = "\033[31m" +TERM_GREEN = "\033[32m" +TERM_YELLOW = "\033[33m" + +logger = logging.getLogger(__name__) +logging.basicConfig(level=logging.INFO) + + [email protected] +class Options: + """Command line options for stub generation.""" + + dlls: list[str] = dataclasses.field(default_factory=list) + indent: int = 4 + files: list[str] = dataclasses.field(default_factory=list) + header_lines: list[str] = dataclasses.field(default_factory=list) Review Comment: fixed -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
