junrushao commented on code in PR #508:
URL: https://github.com/apache/tvm-ffi/pull/508#discussion_r2971877489
##########
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:
good catch!
##########
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:
good catch!
--
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]
