feat: add BFS-based ffi.ReprPrint for unified object repr

Author: junrushao

CMakeLists.txt

@@ -78,6 +78,7 @@ set(_tvm_ffi_extra_objs_sources
     "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/json_writer.cc"
     "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/serialization.cc"
     "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/deep_copy.cc"
+    "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/repr_print.cc"
     "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/reflection_extra.cc"
     "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/module.cc"
     "${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/library_module.cc"

include/tvm/ffi/c_api.h

@@ -870,6 +870,13 @@ typedef enum {
    * being used directly as the default value.
    */
   kTVMFFIFieldFlagBitMaskDefaultFromFactory = 1 << 5,
+  /*!
+   * \brief The field is excluded from repr output.
+   *
+   * When set, the field will not appear in the generic reflection-based repr.
+   * By default this flag is off (meaning the field is included in repr).
+   */
+  kTVMFFIFieldFlagBitMaskReprOff = 1 << 6,
 #ifdef __cplusplus
 };
 #else

include/tvm/ffi/reflection/registry.h

@@ -223,6 +223,33 @@ class AttachFieldFlag : public InfoTrait {
   int32_t flag_;
 };
 
+/*!
+ * \brief Trait that controls whether a field appears in repr output.
+ *
+ * By default, all fields appear in repr. Use `Repr(false)` to exclude a field.
+ */
+class Repr : public InfoTrait {
+ public:
+  /*!
+   * \brief Constructor.
+   * \param show Whether the field should appear in repr output.
+   */
+  explicit Repr(bool show) : show_(show) {}
+
+  /*!
+   * \brief Apply the repr flag to the field info.
+   * \param info The field info.
+   */
+  TVM_FFI_INLINE void Apply(TVMFFIFieldInfo* info) const {
+    if (!show_) {
+      info->flags |= kTVMFFIFieldFlagBitMaskReprOff;
+    }
+  }
+
+ private:
+  bool show_;
+};
+
 /*!
  * \brief Get the byte offset of a class member field.
  *
@@ -493,6 +520,7 @@ struct init {
 namespace type_attr {
 inline constexpr const char* kInit = "__ffi_init__";
 inline constexpr const char* kShallowCopy = "__ffi_shallow_copy__";
+inline constexpr const char* kRepr = "__ffi_repr__";
 }  // namespace type_attr
 
 /*!

python/tvm_ffi/_ffi_api.py

@@ -82,6 +82,7 @@ def ModuleImportModule(_0: Module, _1: Module, /) -> None: ...
     def ModuleInspectSource(_0: Module, _1: str, /) -> str: ...
     def ModuleLoadFromFile(_0: str, /) -> Module: ...
     def ModuleWriteToFile(_0: Module, _1: str, _2: str, /) -> None: ...
+    def ReprPrint(_0: Any, /) -> str: ...
     def Shape(*args: Any) -> Any: ...
     def String(_0: str, /) -> str: ...
     def StructuralHash(_0: Any, _1: bool, _2: bool, /) -> int: ...
@@ -141,6 +142,7 @@ def ToJSONGraphString(_0: Any, _1: Any, /) -> str: ...
     "ModuleInspectSource",
     "ModuleLoadFromFile",
     "ModuleWriteToFile",
+    "ReprPrint",
     "Shape",
     "String",
     "StructuralHash",

python/tvm_ffi/container.py

@@ -197,7 +197,7 @@ def __repr__(self) -> str:
         # exception safety handling for chandle=None
         if self.__chandle__() == 0:
             return type(self).__name__ + "(chandle=None)"
-        return "[" + ", ".join([x.__repr__() for x in self]) + "]"
+        return str(core.__object_repr__(self))  # ty: ignore[unresolved-attribute]
 
     def __contains__(self, value: object) -> bool:
         """Check if the array contains a value."""
@@ -343,7 +343,7 @@ def __repr__(self) -> str:
         """Return a string representation of the list."""
         if self.__chandle__() == 0:
             return type(self).__name__ + "(chandle=None)"
-        return "[" + ", ".join([x.__repr__() for x in self]) + "]"
+        return str(core.__object_repr__(self))  # ty: ignore[unresolved-attribute]
 
     def __contains__(self, value: object) -> bool:
         """Check if the list contains a value."""
@@ -543,4 +543,4 @@ def __repr__(self) -> str:
         # exception safety handling for chandle=None
         if self.__chandle__() == 0:
             return type(self).__name__ + "(chandle=None)"
-        return "{" + ", ".join([f"{k.__repr__()}: {v.__repr__()}" for k, v in self.items()]) + "}"
+        return str(core.__object_repr__(self))  # ty: ignore[unresolved-attribute]

python/tvm_ffi/cython/object.pxi

@@ -26,8 +26,22 @@ def _set_class_object(cls):
     _CLASS_OBJECT = cls
 
 
+_REPR_PRINT = None
+_REPR_PRINT_LOADED = False
+
+
 def __object_repr__(obj: "Object") -> str:
-    """Object repr function that can be overridden by assigning to it"""
+    """Object repr function using ffi.ReprPrint when available."""
+    global _REPR_PRINT, _REPR_PRINT_LOADED
+    if not _REPR_PRINT_LOADED:
+        _REPR_PRINT_LOADED = True
+        _REPR_PRINT = _get_global_func("ffi.ReprPrint", False)
+    if _REPR_PRINT is not None:
+        try:
+            return str(_REPR_PRINT(obj))
+        except Exception:  # noqa: BLE001
+            # Silently fall back: __repr__ must never raise.
+            pass
     return type(obj).__name__ + "(" + str(obj.__ctypes_handle__().value) + ")"
 
 

python/tvm_ffi/dataclasses/_utils.py

@@ -127,46 +127,6 @@ def _get_all_fields(type_info: TypeInfo) -> list[TypeField]:
     return fields
 
 
-def method_repr(type_cls: type, type_info: TypeInfo) -> Callable[..., str]:
-    """Generate a ``__repr__`` method for the dataclass.
-
-    The generated representation includes all fields with ``repr=True`` in
-    the format ``ClassName(field1=value1, field2=value2, ...)``.
-    """
-    # Step 0. Collect all fields from the type hierarchy
-    fields = _get_all_fields(type_info)
-
-    # Step 1. Filter fields that should appear in repr
-    repr_fields: list[str] = []
-    for field in fields:
-        assert field.name is not None
-        assert field.dataclass_field is not None
-        if field.dataclass_field.repr:
-            repr_fields.append(field.name)
-
-    # Step 2. Generate the repr method
-    if not repr_fields:
-        # No fields to show, return a simple class name representation
-        body_lines = [f"return f'{type_cls.__name__}()'"]
-    else:
-        # Build field representations
-        fields_str = ", ".join(
-            f"{field_name}={{self.{field_name}!r}}" for field_name in repr_fields
-        )
-        body_lines = [f"return f'{type_cls.__name__}({fields_str})'"]
-
-    source_lines = ["def __repr__(self) -> str:"]
-    source_lines.extend(f"    {line}" for line in body_lines)
-    source = "\n".join(source_lines)
-
-    # Note: Code generation in this case is guaranteed to be safe,
-    # because the generated code does not contain any untrusted input.
-    namespace: dict[str, Any] = {}
-    exec(source, {}, namespace)
-    __repr__ = namespace["__repr__"]
-    return __repr__
-
-
 def method_init(_type_cls: type, type_info: TypeInfo) -> Callable[..., None]:
     """Generate an ``__init__`` that forwards to the FFI constructor.
 

python/tvm_ffi/dataclasses/c_class.py

@@ -41,7 +41,7 @@
 
 @dataclass_transform(field_specifiers=(field,), kw_only_default=False)
 def c_class(
-    type_key: str, init: bool = True, kw_only: bool = False, repr: bool = True
+    type_key: str, init: bool = True, kw_only: bool = False
 ) -> Callable[[Type[_InputClsType]], Type[_InputClsType]]:  # noqa: UP006
     """(Experimental) Create a dataclass-like proxy for a C++ class registered with TVM FFI.
 
@@ -77,10 +77,6 @@ def c_class(
         ``__init__``. Individual fields can override this by setting
         ``kw_only=False`` in :func:`field`. Additionally, a ``KW_ONLY`` sentinel
         annotation can be used to mark all subsequent fields as keyword-only.
-    repr
-        If ``True`` and the Python class does not define ``__repr__``, a
-        representation method is auto-generated that includes all fields with
-        ``repr=True``.
 
     Returns
     -------
@@ -128,9 +124,8 @@ class MyClass:
     """
 
     def decorator(super_type_cls: Type[_InputClsType]) -> Type[_InputClsType]:  # noqa: UP006
-        nonlocal init, repr
+        nonlocal init
         init = init and "__init__" not in super_type_cls.__dict__
-        repr = repr and "__repr__" not in super_type_cls.__dict__
         # Step 1. Retrieve `type_info` from registry
         type_info: TypeInfo = _lookup_or_register_type_info_from_type_key(type_key)
         assert type_info.parent_type_info is not None
@@ -146,11 +141,10 @@ def decorator(super_type_cls: Type[_InputClsType]) -> Type[_InputClsType]:  # no
             )
         # Step 3. Create the proxy class with the fields as properties
         fn_init = _utils.method_init(super_type_cls, type_info) if init else None
-        fn_repr = _utils.method_repr(super_type_cls, type_info) if repr else None
         type_cls: Type[_InputClsType] = _utils.type_info_to_cls(  # noqa: UP006
             type_info=type_info,
             cls=super_type_cls,
-            methods={"__init__": fn_init, "__repr__": fn_repr},
+            methods={"__init__": fn_init},
         )
         _set_type_cls(type_info, type_cls)
         # Step 4. Set up __copy__, __deepcopy__, __replace__

python/tvm_ffi/dataclasses/field.py

@@ -47,22 +47,20 @@ class Field:
     way the decorator understands.
     """
 
-    __slots__ = ("default_factory", "init", "kw_only", "name", "repr")
+    __slots__ = ("default_factory", "init", "kw_only", "name")
 
     def __init__(
         self,
         *,
         name: str | None = None,
         default_factory: Callable[[], _FieldValue] | _MISSING_TYPE = MISSING,
         init: bool = True,
-        repr: bool = True,
         kw_only: bool | _MISSING_TYPE = MISSING,
     ) -> None:
         """Do not call directly; use :func:`field` instead."""
         self.name = name
         self.default_factory = default_factory
         self.init = init
-        self.repr = repr
         self.kw_only = kw_only
 
 
@@ -71,7 +69,6 @@ def field(
     default: _FieldValue | _MISSING_TYPE = MISSING,
     default_factory: Callable[[], _FieldValue] | _MISSING_TYPE = MISSING,
     init: bool = True,
-    repr: bool = True,
     kw_only: bool | _MISSING_TYPE = MISSING,
 ) -> _FieldValue:
     """(Experimental) Declare a dataclass-style field on a :func:`c_class` proxy.
@@ -94,9 +91,6 @@ def field(
     init
         If ``True`` the field is included in the generated ``__init__``.
         If ``False`` the field is omitted from input arguments of ``__init__``.
-    repr
-        If ``True`` the field is included in the generated ``__repr__``.
-        If ``False`` the field is omitted from the ``__repr__`` output.
     kw_only
         If ``True``, the field is a keyword-only argument in ``__init__``.
         If ``MISSING``, inherits from the class-level ``kw_only`` setting or
@@ -158,13 +152,11 @@ class PyBase:
         raise ValueError("Cannot specify both `default` and `default_factory`")
     if not isinstance(init, bool):
         raise TypeError("`init` must be a bool")
-    if not isinstance(repr, bool):
-        raise TypeError("`repr` must be a bool")
     if kw_only is not MISSING and not isinstance(kw_only, bool):
         raise TypeError(f"`kw_only` must be a bool, got {type(kw_only).__name__!r}")
     if default is not MISSING:
         default_factory = _make_default_factory(default)
-    ret = Field(default_factory=default_factory, init=init, repr=repr, kw_only=kw_only)
+    ret = Field(default_factory=default_factory, init=init, kw_only=kw_only)
     return cast(_FieldValue, ret)