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 a39cddaa Improve typing of overloaded methods in stubgen (#678)
a39cddaa is described below
commit a39cddaace892987787449632a9c9b39725c1bfa
Author: Matthew Brookhart <[email protected]>
AuthorDate: Mon Jul 20 18:01:46 2026 -0600
Improve typing of overloaded methods in stubgen (#678)
---
include/tvm/ffi/reflection/overload.h | 30 +++++++++++-----
python/tvm_ffi/stub/python_generator/codegen.py | 2 ++
python/tvm_ffi/stub/python_generator/utils.py | 47 ++++++++++++++-----------
python/tvm_ffi/stub/utils.py | 10 ++++++
tests/cpp/test_overload.cc | 15 ++++++++
tests/python/test_stubgen.py | 39 +++++++++++++++++---
6 files changed, 110 insertions(+), 33 deletions(-)
diff --git a/include/tvm/ffi/reflection/overload.h
b/include/tvm/ffi/reflection/overload.h
index 04c21fb5..c9193c40 100644
--- a/include/tvm/ffi/reflection/overload.h
+++ b/include/tvm/ffi/reflection/overload.h
@@ -493,18 +493,25 @@ class OverloadObjectDef : private ObjectDef<Class> {
auto method_name = std::string(type_key_) + "." + name;
- // if an overload method exists, register to existing overload function
+ // If an overload method exists, add the callable to the existing runtime
+ // dispatcher but still publish a TypeMethod entry for this signature.
+ // Stub generators consume TypeInfo::methods, so omitting later entries
+ // made an overloaded C++ method look like a single Python signature.
+ Function method;
if (const auto overload_it = registered_fields_.find(name);
overload_it != registered_fields_.end()) {
- ::tvm::ffi::details::OverloadBase* overload_ptr = overload_it->second;
- return overload_ptr->Register(NewOverload(std::move(method_name),
std::forward<Func>(func)));
+ overload_it->second.overload->Register(
+ NewOverload(std::move(method_name), std::forward<Func>(func)));
+ method = overload_it->second.method;
+ } else {
+ // First registration creates the runtime overload dispatcher. Keep the
+ // Function alive so every reflected signature can reference it.
+ auto [new_method, overload_ptr] =
+ GetOverloadMethod(std::move(method_name), std::forward<Func>(func));
+ method = std::move(new_method);
+ registered_fields_.try_emplace(name, RegisteredMethod{overload_ptr,
method});
}
- // first time registering overload method
- auto [method, overload_ptr] =
- GetOverloadMethod(std::move(method_name), std::forward<Func>(func));
- registered_fields_.try_emplace(name, overload_ptr);
-
info.method = AnyView(method).CopyToTVMFFIAny();
info.metadata_.emplace_back("type_schema", FuncInfo::TypeSchema());
// apply method info traits
@@ -514,7 +521,12 @@ class OverloadObjectDef : private ObjectDef<Class> {
TVM_FFI_CHECK_SAFE_CALL(TVMFFITypeRegisterMethod(type_index_, &info));
}
- std::unordered_map<std::string, ::tvm::ffi::details::OverloadBase*>
registered_fields_;
+ struct RegisteredMethod {
+ ::tvm::ffi::details::OverloadBase* overload;
+ Function method;
+ };
+
+ std::unordered_map<std::string, RegisteredMethod> registered_fields_;
};
} // namespace reflection
diff --git a/python/tvm_ffi/stub/python_generator/codegen.py
b/python/tvm_ffi/stub/python_generator/codegen.py
index 0f813464..2669231d 100644
--- a/python/tvm_ffi/stub/python_generator/codegen.py
+++ b/python/tvm_ffi/stub/python_generator/codegen.py
@@ -183,6 +183,8 @@ def generate_python_object(
info = obj_info
method_names = {m.schema.name.rsplit(".", 1)[-1] for m in info.methods}
fn_ty_map = _type_suffix_and_record(ty_map, imports,
func_names=method_names)
+ if info.has_overloaded_methods():
+ imports.append(ImportItem("typing.overload", type_checking_only=True))
input_fn_ty_map = _type_suffix_and_record(
_make_input_ty_map(ty_map), imports, func_names=method_names
)
diff --git a/python/tvm_ffi/stub/python_generator/utils.py
b/python/tvm_ffi/stub/python_generator/utils.py
index 37eff8f1..bf19b367 100644
--- a/python/tvm_ffi/stub/python_generator/utils.py
+++ b/python/tvm_ffi/stub/python_generator/utils.py
@@ -155,16 +155,29 @@ def render_object_methods(
input_ty_map = ty_map
indent_str = " " * indent
ret = []
+ groups: list[list[FuncInfo]] = []
+ seen: dict[tuple[str, bool], list[FuncInfo]] = {}
for method in info.methods:
- func_name = method.schema.name.rsplit(".", 1)[-1]
- if func_name == "__ffi_init__":
- # __ffi_init__ is installed as an instance method (self, *args,
**kwargs) -> None
- # by _install_ffi_init_attr, regardless of the C++ static
registration.
- ret.append(_render_ffi_init_from_method(method, ty_map, indent,
input_ty_map))
- continue
- if not method.is_member:
- ret.append(f"{indent_str}@staticmethod")
- ret.append(render_func_signature(method, ty_map, indent, input_ty_map))
+ key = (method.schema.name, method.is_member)
+ if key not in seen:
+ group: list[FuncInfo] = []
+ groups.append(group)
+ seen[key] = group
+ seen[key].append(method)
+
+ for group in groups:
+ for candidate in group:
+ func_name = candidate.schema.name.rsplit(".", 1)[-1]
+ if len(group) > 1:
+ ret.append(f"{indent_str}@overload")
+ if func_name == "__ffi_init__":
+ # __ffi_init__ is installed as an instance method (self,
*args, **kwargs) -> None
+ # by _install_ffi_init_attr, regardless of the C++ static
registration.
+ ret.append(_render_ffi_init_from_method(candidate, ty_map,
indent, input_ty_map))
+ continue
+ if not candidate.is_member:
+ ret.append(f"{indent_str}@staticmethod")
+ ret.append(render_func_signature(candidate, ty_map, indent,
input_ty_map))
return ret
@@ -179,20 +192,17 @@ def _render_ffi_init_from_method(
input_ty_map = ty_map
indent_str = " " * indent
schema = method.schema
- # Subclass __ffi_init__ signatures legitimately differ from the parent
- # (different fields -> different constructor params), so suppress LSP.
- ignore = " # ty: ignore[invalid-method-override]"
if schema.origin != "Callable" or not schema.args:
ty_map("Any")
- return f"{indent_str}def __ffi_init__(self, *args: Any) -> None:
...{ignore}"
+ return f"{indent_str}def __ffi_init__(self, *args: Any) -> None: ..."
# schema.args[0] is return type, schema.args[1:] are param types.
parts: list[str] = []
for i, arg in enumerate(schema.args[1:]):
parts.append(f"_{i}: {arg.input_repr(input_ty_map)}")
if parts:
params = ", ".join(parts)
- return f"{indent_str}def __ffi_init__(self, {params}, /) -> None:
...{ignore}"
- return f"{indent_str}def __ffi_init__(self) -> None: ...{ignore}"
+ return f"{indent_str}def __ffi_init__(self, {params}, /) -> None: ..."
+ return f"{indent_str}def __ffi_init__(self) -> None: ..."
def render_object_ffi_init(
@@ -281,10 +291,7 @@ def _render_ffi_init_from_fields(
) -> list[str]:
"""Render ``__ffi_init__`` stub from field metadata for auto-generated
init."""
indent_str = " " * indent
- # Subclass __ffi_init__ signatures legitimately differ from the parent
- # (different fields -> different constructor params), so suppress LSP.
- ignore = " # ty: ignore[invalid-method-override]"
params = _format_field_params(info, ty_map, input_ty_map)
if params:
- return [f"{indent_str}def __ffi_init__(self, {params}) -> None:
...{ignore}"]
- return [f"{indent_str}def __ffi_init__(self) -> None: ...{ignore}"]
+ return [f"{indent_str}def __ffi_init__(self, {params}) -> None: ..."]
+ return [f"{indent_str}def __ffi_init__(self) -> None: ..."]
diff --git a/python/tvm_ffi/stub/utils.py b/python/tvm_ffi/stub/utils.py
index 172fac33..07f40f8d 100644
--- a/python/tvm_ffi/stub/utils.py
+++ b/python/tvm_ffi/stub/utils.py
@@ -130,6 +130,16 @@ class ObjectInfo:
init_fields: list[InitFieldInfo] = dataclasses.field(default_factory=list)
has_init: bool = False
+ def has_overloaded_methods(self) -> bool:
+ """Return whether reflection exposed multiple signatures for a
method."""
+ seen: set[tuple[str, bool]] = set()
+ for method in self.methods:
+ key = (method.schema.name, method.is_member)
+ if key in seen:
+ return True
+ seen.add(key)
+ return False
+
@staticmethod
def from_type_info(type_info: TypeInfo) -> ObjectInfo:
"""Construct an `ObjectInfo` from a `TypeInfo` instance."""
diff --git a/tests/cpp/test_overload.cc b/tests/cpp/test_overload.cc
index 7dfb9c70..7c2f62c7 100644
--- a/tests/cpp/test_overload.cc
+++ b/tests/cpp/test_overload.cc
@@ -26,6 +26,8 @@
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/string.h>
+#include <string>
+
namespace {
using namespace tvm::ffi;
@@ -92,4 +94,17 @@ TEST(Reflection, CallOverloadedStaticMethod) {
EXPECT_EQ(res_b.as<float>(), 2.0f);
}
+TEST(Reflection, RecordsEveryOverloadedMethodSignature) {
+ const TypeInfo* info =
TVMFFIGetTypeInfo(TestOverloadObj::RuntimeTypeIndex());
+ int static_overload_count = 0;
+ for (int32_t index = 0; index < info->num_methods; ++index) {
+ const TVMFFIMethodInfo& method = info->methods[index];
+ if (std::string(method.name.data, method.name.size) == "add_one_static") {
+ ++static_overload_count;
+ EXPECT_GT(method.metadata.size, 0);
+ }
+ }
+ EXPECT_EQ(static_overload_count, 2);
+}
+
} // namespace
diff --git a/tests/python/test_stubgen.py b/tests/python/test_stubgen.py
index 277c725d..f94d2418 100644
--- a/tests/python/test_stubgen.py
+++ b/tests/python/test_stubgen.py
@@ -257,6 +257,40 @@ def test_objectinfo_gen_fields_and_methods() -> None:
]
+def test_objectinfo_gen_overloaded_static_methods() -> None:
+ info = ObjectInfo(
+ fields=[],
+ methods=[
+ FuncInfo.from_schema(
+ "demo.Factory.create",
+ TypeSchema("Callable", (TypeSchema("demo.Factory"),
TypeSchema("int"))),
+ is_member=False,
+ ),
+ FuncInfo.from_schema(
+ "demo.Factory.reset",
+ TypeSchema("Callable", (TypeSchema("None"),)),
+ is_member=False,
+ ),
+ FuncInfo.from_schema(
+ "demo.Factory.create",
+ TypeSchema("Callable", (TypeSchema("demo.Factory"),
TypeSchema("str"))),
+ is_member=False,
+ ),
+ ],
+ )
+ assert info.has_overloaded_methods()
+ assert render_object_methods(info, _identity_ty_map, indent=2) == [
+ " @overload",
+ " @staticmethod",
+ " def create(_0: int, /) -> demo.Factory: ...",
+ " @overload",
+ " @staticmethod",
+ " def create(_0: str, /) -> demo.Factory: ...",
+ " @staticmethod",
+ " def reset() -> None: ...",
+ ]
+
+
def test_type_schema_container_origins() -> None:
"""Test that Array/List/Map/Dict origins are distinct and validated
correctly."""
# Array and List: 0 or 1 arg, default to (Any,)
@@ -400,10 +434,7 @@ def test_objectinfo_gen_init_uses_input_annotations() ->
None:
]
assert render_object_ffi_init(
info, _type_suffix, indent=0, input_ty_map=_input_type_suffix
- ) == [
- "def __ffi_init__(self, items: Sequence[int]) -> None: ... # ty: "
- "ignore[invalid-method-override]"
- ]
+ ) == ["def __ffi_init__(self, items: Sequence[int]) -> None: ..."]
def test_py_class_method_metadata_renders_stub_signature() -> None:
