[PY] Distinguish FFI input and output annotations (#621)

## Summary
- Add separate input/output rendering APIs for `TypeSchema` so stubgen
can widen callable and constructor inputs without weakening returned or
field types.
- Add a `__ffi_convert_type_schema__` TypeAttr hook for object
conversions and wire `py_class`/`c_class` dataclass transform converter
metadata.
- Generate `type_schema` metadata for Python-defined `@method`
TypeMethods so reflection-driven stubgen can render their signatures.
- Cover container asymmetry, object convert schema recursion, stubgen
signatures, partial type maps, `@method` metadata, and dataclass
metadata with regression tests.
- Update the PR CI workflow to avoid Apache-disallowed third-party setup
actions and run pre-commit without syncing the project during lint.

## Validation
- pre-commit hooks during commits, including ruff, ty, clang-format,
cython-lint, CMake lint/format, and workflow YAML checks
- `python -m pytest
tests/python/test_stubgen.py::test_py_class_method_metadata_renders_stub_signature
tests/python/test_typed_method.py -q`
- `python -m pytest tests/python/test_stubgen.py -q`
- `CUDA_VISIBLE_DEVICES="" python -m pytest tests/python -q` (`2322
passed, 18 skipped, 2 xfailed`)
- Manual output dumped under
`tmp/2026-06-14-update-ffi/manual_check.out`
- GitHub CI is green for lint, docs, Ubuntu, Ubuntu ARM, macOS, and
Windows

Author: Ubospica

CMakeLists.txt

@@ -247,10 +247,13 @@ if (TVM_FFI_BUILD_PYTHON_MODULE)
   set(_cython_sources
       ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/core.pyx
       ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/base.pxi
+      ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/type_info.pxi
       ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/device.pxi
       ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/dtype.pxi
       ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/error.pxi
       ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/function.pxi
+      ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/pycallback.pxi
+      ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/pyclass_type_converter.pxi
       ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/tensor.pxi
       ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/object.pxi
       ${CMAKE_CURRENT_SOURCE_DIR}/python/tvm_ffi/cython/string.pxi

cmake/Utils/AddGoogleTest.cmake

@@ -91,7 +91,7 @@ macro (TVM_FFI_ADD_GTEST target_name)
   target_link_libraries(${target_name} PRIVATE gtest_main)
   gtest_discover_tests(
     ${target_name}
-    WORKING_DIRECTORY ${PROJECT_SOURCE_DIR} TEST_DISCOVERY_TIMEOUT 600 DISCOVERY_MODE PRE_TEST
+    WORKING_DIRECTORY ${PROJECT_SOURCE_DIR} DISCOVERY_TIMEOUT 600 DISCOVERY_MODE PRE_TEST
     PROPERTIES VS_DEBUGGER_WORKING_DIRECTORY "${PROJECT_SOURCE_DIR}"
   )
   set_target_properties(${target_name} PROPERTIES FOLDER tests)

include/tvm/ffi/reflection/accessor.h

@@ -350,6 +350,14 @@ inline constexpr const char* kInit = "__ffi_init__";
  * Signature: ``(AnyView src) -> TSelf``, where ``TSelf`` is a subclass of ObjectRef.
  */
 inline constexpr const char* kConvert = "__ffi_convert__";
+/*!
+ * \brief Type schema accepted by ``kConvert`` before it returns ``TSelf``.
+ *
+ * Stored as a JSON type schema string or schema-like object.  Python stub
+ * generation uses this attribute to render widened input annotations while
+ * keeping output annotations precise.
+ */
+inline constexpr const char* kConvertTypeSchema = "__ffi_convert_type_schema__";
 /*!
  * \brief Shallow-copy factory.
  *

python/tvm_ffi/core.pyi

@@ -274,6 +274,8 @@ class TypeSchema:
     @staticmethod
     def from_annotation(annotation: object) -> TypeSchema: ...
     def repr(self, ty_map: Callable[[str], str] | None = None) -> str: ...
+    def input_repr(self, ty_map: Callable[[str], str] | None = None) -> str: ...
+    def output_repr(self, ty_map: Callable[[str], str] | None = None) -> str: ...
     def check_value(self, value: object) -> None: ...
     def convert(self, value: object) -> CAny: ...
     def to_json(self) -> dict[str, Any]: ...

python/tvm_ffi/cython/type_info.pxi

@@ -116,6 +116,8 @@ _TYPE_SCHEMA_ORIGIN_CONVERTER = {
     "ObjectRValueRef": "Object",
 }
 
+cdef str _TYPE_ATTR_FFI_CONVERT_TYPE_SCHEMA = "__ffi_convert_type_schema__"
+
 # Sentinel for structural types (Optional, Union) that have no single type_index
 _ORIGIN_TYPE_INDEX_STRUCTURAL = -2
 # Sentinel for unknown/unresolved origins
@@ -217,6 +219,19 @@ class TypeSchema:
     def __repr__(self) -> str:
         return self.repr(ty_map=None)
 
+    @staticmethod
+    def _from_maybe_json(raw: object) -> "TypeSchema":
+        """Construct a TypeSchema from a TypeSchema, JSON string, or JSON dict."""
+        if isinstance(raw, TypeSchema):
+            return raw
+        if isinstance(raw, str):
+            return TypeSchema.from_json_str(raw)
+        if isinstance(raw, dict):
+            return TypeSchema.from_json_obj(raw)
+        raise TypeError(
+            f"expected TypeSchema, JSON string, or JSON dict, got {type(raw).__name__}"
+        )
+
     @staticmethod
     def from_json_obj(obj: dict[str, Any]) -> "TypeSchema":
         """Construct a :class:`TypeSchema` from a parsed JSON object.
@@ -517,12 +532,40 @@ class TypeSchema:
             assert s.repr() == "Array[int]"
 
         """
-        if ty_map is None:
-            origin = self.origin
-        else:
-            origin = ty_map(self.origin)
+        return self.output_repr(ty_map)
+
+    def input_repr(self, ty_map: "Optional[Callable[[str], str]]" = None) -> str:
+        """Render the Python input annotation accepted by this schema."""
+        return self._repr_impl(ty_map, input_mode=True, expanded_convert_types=frozenset())
+
+    def output_repr(self, ty_map: "Optional[Callable[[str], str]]" = None) -> str:
+        """Render the precise Python output annotation produced by this schema."""
+        return self._repr_impl(ty_map, input_mode=False, expanded_convert_types=frozenset())
+
+    def _repr_impl(
+        self,
+        ty_map: "Optional[Callable[[str], str]]",
+        input_mode: bool,
+        expanded_convert_types: "frozenset[int]",
+    ) -> str:
+        if input_mode and self.origin_type_index >= kTVMFFIStaticObjectBegin:
+            if self.origin_type_index not in expanded_convert_types:
+                raw_convert_schema = _lookup_type_attr(
+                    self.origin_type_index, _TYPE_ATTR_FFI_CONVERT_TYPE_SCHEMA
+                )
+                if raw_convert_schema is not None:
+                    return TypeSchema._from_maybe_json(raw_convert_schema)._repr_impl(
+                        ty_map,
+                        input_mode=True,
+                        expanded_convert_types=expanded_convert_types | {self.origin_type_index},
+                    )
+
+        origin = self.origin if ty_map is None else ty_map(self.origin)
         schema_args = self.args
-        args = [i.repr(ty_map) for i in (() if schema_args is None else schema_args)]
+        args = [
+            i._repr_impl(ty_map, input_mode, expanded_convert_types)
+            for i in (() if schema_args is None else schema_args)
+        ]
         if origin == "Union":
             return " | ".join(args)
         elif origin == "Optional":
@@ -1129,8 +1172,8 @@ cdef _register_py_methods(int32_t type_index, list py_methods, frozenset type_at
     ----------
     type_index : int
         The runtime type index of the type.
-    py_methods : list[tuple[str, Any, bool]]
-        Each entry is ``(name, value, is_static)``.
+    py_methods : list[tuple[str, Any, bool, str | None]]
+        Each entry is ``(name, value, is_static, metadata_json)``.
     type_attr_names : frozenset[str]
         Names to register as TypeAttrColumn instead of TypeMethod.
     """
@@ -1139,11 +1182,12 @@ cdef _register_py_methods(int32_t type_index, list py_methods, frozenset type_at
     cdef TVMFFIAny sentinel_any
     cdef int c_api_ret_code
     cdef ByteArrayArg name_arg
+    cdef ByteArrayArg metadata_arg
 
     sentinel_any.type_index = kTVMFFINone
     sentinel_any.v_int64 = 0
 
-    for name, func, is_static in py_methods:
+    for name, func, is_static, metadata_json in py_methods:
         func_any.type_index = kTVMFFINone
         func_any.v_int64 = 0
         try:
@@ -1169,8 +1213,13 @@ cdef _register_py_methods(int32_t type_index, list py_methods, frozenset type_at
                 method_info.doc.size = 0
                 method_info.flags = kTVMFFIFieldFlagBitMaskIsStaticMethod if is_static else 0
                 method_info.method = func_any
-                method_info.metadata.data = NULL
-                method_info.metadata.size = 0
+                if metadata_json is not None:
+                    metadata_bytes = c_str(metadata_json)
+                    metadata_arg = ByteArrayArg(metadata_bytes)
+                    method_info.metadata = metadata_arg.cdata
+                else:
+                    method_info.metadata.data = NULL
+                    method_info.metadata.size = 0
                 CHECK_CALL(TVMFFITypeRegisterMethod(type_index, &method_info))
         finally:
             if func_any.type_index >= kTVMFFIStaticObjectBegin and func_any.v_obj != NULL:

python/tvm_ffi/dataclasses/c_class.py

@@ -24,7 +24,7 @@
 
 from typing_extensions import dataclass_transform
 
-from .field import Field
+from .field import Field, _field_converter, field
 
 _T = TypeVar("_T", bound=type)
 
@@ -60,7 +60,12 @@ def _attach_field_objects(cls: type, type_info: Any) -> None:
         tf.dataclass_field = f
 
 
-@dataclass_transform(eq_default=False, order_default=False)
+@dataclass_transform(
+    eq_default=False,
+    order_default=False,
+    field_specifiers=(field, Field),
+    converter=_field_converter,
+)
 def c_class(
     type_key: str,
     *,

python/tvm_ffi/dataclasses/field.py

@@ -37,6 +37,11 @@ class KW_ONLY:
         """Sentinel type: annotations after ``_: KW_ONLY`` are keyword-only."""
 
 
+def _field_converter(value: Any) -> Any:
+    """Static-analysis marker for fields whose values are converted by FFI."""
+    return value
+
+
 class Field:
     """Descriptor for a single field in a Python-defined TVM-FFI type.
 
@@ -102,12 +107,16 @@ class Field:
           parameters that reference outer-scope vars.
     doc : str | None
         Optional docstring for the field.
+    converter : Callable[[Any], Any]
+        Static-analysis marker for field conversion. Runtime conversion is
+        still handled by the FFI type converter.
 
     """
 
     __slots__ = (
         "_ty_schema",
         "compare",
+        "converter",
         "default",
         "default_factory",
         "doc",
@@ -130,6 +139,7 @@ class Field:
     repr: bool
     hash: bool | None
     compare: bool
+    converter: Callable[[Any], Any]
     kw_only: bool | None
     structural_eq: str | None
     doc: str | None
@@ -158,6 +168,7 @@ def __init__(  # noqa: PLR0913
         kw_only: bool | None = False,
         structural_eq: str | None = None,
         doc: str | None = None,
+        converter: Callable[[Any], Any] = _field_converter,
     ) -> None:
         # MISSING means "parameter not provided".
         # An explicit None from the user fails the callable() check,
@@ -185,12 +196,13 @@ def __init__(  # noqa: PLR0913
         self.repr = repr
         self.hash = hash
         self.compare = compare
+        self.converter = converter
         self.kw_only = kw_only
         self.structural_eq = structural_eq
         self.doc = doc
 
 
-def field(
+def field(  # noqa: PLR0913
     *,
     default: object = MISSING,
     default_factory: Callable[[], object] | None = MISSING,  # type: ignore[assignment]
@@ -202,6 +214,7 @@ def field(
     kw_only: bool | None = None,
     structural_eq: str | None = None,
     doc: str | None = None,
+    converter: Callable[[Any], Any] = _field_converter,
 ) -> Any:
     """Customize a field in a ``@py_class``-decorated class.
 
@@ -248,6 +261,9 @@ def field(
         binding.
     doc
         Optional docstring for the field.
+    converter
+        Static-analysis marker for field conversion. Runtime conversion is
+        still handled by the FFI type converter.
 
     Returns
     -------
@@ -282,4 +298,5 @@ class MyFunc(Object):
         kw_only=kw_only,
         structural_eq=structural_eq,
         doc=doc,
+        converter=converter,
     )