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 0ee6444 Minor improvement on type hint of `register_object` (#212)
0ee6444 is described below
commit 0ee644421cae936c459a3e0ff41863cf07259647
Author: Kexing Zhou <[email protected]>
AuthorDate: Sat Nov 1 07:25:51 2025 +0800
Minor improvement on type hint of `register_object` (#212)
Update type hint for register_object function.
This pull request introduces a minor type hint improvement to the
`register_object` function in `python/tvm_ffi/registry.py`. The change
enhances type safety by introducing a type variable for the function's
return type.
* Improved type hinting for `register_object` by adding a `_T` type
variable, ensuring the decorator preserves the original type of the
class it decorates.
Before, the decorated object is treated as a value, and cause type lint
errors
<img width="592" height="171" alt="image"
src="https://github.com/user-attachments/assets/a5cd248a-c48e-4755-8706-c836993ee6f0";
/>
After, the object is treated as a type
<img width="596" height="183" alt="image"
src="https://github.com/user-attachments/assets/4b0bd484-7e6a-4830-b10a-bec8f90da9de";
/>
---------
Co-authored-by: Junru Shao <[email protected]>
---
python/tvm_ffi/registry.py | 13 ++++++++-----
1 file changed, 8 insertions(+), 5 deletions(-)
diff --git a/python/tvm_ffi/registry.py b/python/tvm_ffi/registry.py
index d149488..ce1d15f 100644
--- a/python/tvm_ffi/registry.py
+++ b/python/tvm_ffi/registry.py
@@ -20,7 +20,7 @@ from __future__ import annotations
import json
import sys
-from typing import Any, Callable, Literal, overload
+from typing import Any, Callable, Literal, TypeVar, overload
from . import core
from .core import TypeInfo
@@ -29,7 +29,10 @@ from .core import TypeInfo
_SKIP_UNKNOWN_OBJECTS = False
-def register_object(type_key: str | type | None = None) -> Callable[[type],
type] | type:
+_T = TypeVar("_T", bound=type)
+
+
+def register_object(type_key: str | None = None) -> Callable[[_T], _T]:
"""Register object type.
Parameters
@@ -51,7 +54,7 @@ def register_object(type_key: str | type | None = None) ->
Callable[[type], type
"""
- def _register(cls: type, object_name: str) -> type:
+ def _register(cls: _T, object_name: str) -> _T:
"""Register the object type with the FFI core."""
type_index = core._object_type_key_to_index(object_name)
if type_index is None:
@@ -65,12 +68,12 @@ def register_object(type_key: str | type | None = None) ->
Callable[[type], type
if isinstance(type_key, str):
- def _decorator_with_name(cls: type) -> type:
+ def _decorator_with_name(cls: _T) -> _T:
return _register(cls, type_key)
return _decorator_with_name
- def _decorator_default(cls: type) -> type:
+ def _decorator_default(cls: _T) -> _T:
return _register(cls, cls.__name__)
if type_key is None:
