gemini-code-assist[bot] commented on code in PR #508:
URL: https://github.com/apache/tvm-ffi/pull/508#discussion_r2971866468
##########
python/tvm_ffi/cython/type_info.pxi:
##########
@@ -996,6 +1014,64 @@ def _register_fields(type_info, fields):
return type_fields
+cdef _register_py_methods(int32_t type_index, list py_methods):
+ """Register user-defined dunder methods as both TypeMethod and TypeAttr.
+
+ For each method in *py_methods*:
+ 1. Convert the Python callable to a ``TVMFFIAny`` (``ffi::Function``).
+ 2. Call ``TVMFFITypeRegisterMethod`` so the method appears in the
+ type's reflection metadata (``TypeInfo.methods``).
+ 3. Ensure the type-attribute column exists (sentinel call with
+ ``type_index = kTVMFFINone``), then call ``TVMFFITypeRegisterAttr``
+ so the C++ runtime dispatch can find the hook.
+
+ Parameters
+ ----------
+ type_index : int
+ The runtime type index of the type.
+ py_methods : list[tuple[str, callable, bool]]
+ Each entry is ``(name, func, is_static)``.
+ """
+ cdef TVMFFIMethodInfo method_info
+ cdef TVMFFIAny func_any
+ cdef int c_api_ret_code
+ cdef ByteArrayArg name_arg
+ cdef TVMFFIAny sentinel_any
+
+ sentinel_any.type_index = kTVMFFINone
+ sentinel_any.v_int64 = 0
+
+ for name, func, is_static in py_methods:
+ name_bytes = c_str(name)
+ name_arg = ByteArrayArg(name_bytes)
+
+ # Convert Python callable -> TVMFFIAny (creates a FunctionObj)
+ func_any.type_index = kTVMFFINone
+ func_any.v_int64 = 0
+ TVMFFIPyPyObjectToFFIAny(
+ TVMFFIPyArgSetterFactory_,
+ <PyObject*>func,
+ &func_any,
+ &c_api_ret_code,
+ )
+ CHECK_CALL(c_api_ret_code)
+
+ # 1. Register as TypeMethod
+ method_info.name = name_arg.cdata
+ method_info.doc.data = NULL
+ 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
+ CHECK_CALL(TVMFFITypeRegisterMethod(type_index, &method_info))
+
+ # 2. Ensure type-attr column exists (sentinel: kTVMFFINone)
+ CHECK_CALL(TVMFFITypeRegisterAttr(kTVMFFINone, &name_arg.cdata,
&sentinel_any))
+ # 3. Register as TypeAttr
+ CHECK_CALL(TVMFFITypeRegisterAttr(type_index, &name_arg.cdata,
&func_any))
Review Comment:
The `func_any` variable, which holds an owned reference to a `FunctionObj`,
is declared outside the loop and reused. This causes a reference leak in each
iteration because the old object reference is overwritten without being
decremented. To fix this, the `func_any` variable should be managed within the
loop's scope, and its reference must be released after use. A `try...finally`
block is a robust way to ensure the reference is always decremented, preventing
memory leaks even if errors occur.
```
cdef TVMFFIMethodInfo method_info
cdef TVMFFIAny sentinel_any
sentinel_any.type_index = kTVMFFINone
sentinel_any.v_int64 = 0
for name, func, is_static in py_methods:
cdef TVMFFIAny func_any
func_any.type_index = kTVMFFINone
func_any.v_int64 = 0
cdef int c_api_ret_code
try:
name_bytes = c_str(name)
name_arg = ByteArrayArg(name_bytes)
# Convert Python callable -> TVMFFIAny (creates a FunctionObj)
TVMFFIPyPyObjectToFFIAny(
TVMFFIPyArgSetterFactory_,
<PyObject*>func,
&func_any,
&c_api_ret_code,
)
CHECK_CALL(c_api_ret_code)
# 1. Register as TypeMethod
method_info.name = name_arg.cdata
method_info.doc.data = NULL
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
CHECK_CALL(TVMFFITypeRegisterMethod(type_index, &method_info))
# 2. Ensure type-attr column exists (sentinel: kTVMFFINone)
CHECK_CALL(TVMFFITypeRegisterAttr(kTVMFFINone, &name_arg.cdata,
&sentinel_any))
# 3. Register as TypeAttr
CHECK_CALL(TVMFFITypeRegisterAttr(type_index, &name_arg.cdata,
&func_any))
finally:
if func_any.type_index >= kTVMFFIStaticObjectBegin and
func_any.v_obj != NULL:
TVMFFIObjectDecRef(<TVMFFIObjectHandle>func_any.v_obj)
```
##########
python/tvm_ffi/dataclasses/py_class.py:
##########
@@ -202,6 +202,31 @@ def _collect_own_fields(
return fields
+def _collect_py_methods(cls: type) -> list[tuple[str, Any, bool]] | None:
+ """Extract recognized FFI dunder methods from the class body.
+
+ Only names listed in :data:`_FFI_RECOGNIZED_METHODS` are collected.
+
+ Returns a list of ``(name, func, is_static)`` tuples, or ``None``
+ if no eligible methods were found.
+ """
+ methods: list[tuple[str, Any, bool]] = []
+ for name in list(cls.__dict__):
+ if name not in _FFI_RECOGNIZED_METHODS:
+ continue
+ value = cls.__dict__[name]
Review Comment:
To improve efficiency and readability, you can iterate directly over the
dictionary's items using `cls.__dict__.items()`. This avoids creating a
temporary list of keys and performing a second dictionary lookup for the value
inside the loop.
```suggestion
for name, value in cls.__dict__.items():
if name not in _FFI_RECOGNIZED_METHODS:
continue
```
--
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]
