Skip to content

Commit 997cbf1

Browse files
committed
feat: add DFS-based ffi.ReprPrint for unified object repr
- Single C++ ffi.ReprPrint function handles all types - DFS with 3-state tracking (NotVisited/InProgress/Done): - DAGs: memoized repr returned in full on re-encounter - Cycles: detected via InProgress state, shown as ... - Addresses hidden by default; set TVM_FFI_REPR_WITH_ADDR=1 to show - Per-field Repr(false) to exclude fields from repr output - Built-in repr for String, Bytes, Tensor, Shape, Array, List, Map - All Python __repr__ methods delegate to this function
1 parent 3b26a09 commit 997cbf1

14 files changed

Lines changed: 1029 additions & 103 deletions

File tree

CMakeLists.txt

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,7 @@ set(_tvm_ffi_extra_objs_sources
7878
"${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/json_writer.cc"
7979
"${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/serialization.cc"
8080
"${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/deep_copy.cc"
81+
"${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/repr_print.cc"
8182
"${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/reflection_extra.cc"
8283
"${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/module.cc"
8384
"${CMAKE_CURRENT_SOURCE_DIR}/src/ffi/extra/library_module.cc"

include/tvm/ffi/c_api.h

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -870,6 +870,13 @@ typedef enum {
870870
* being used directly as the default value.
871871
*/
872872
kTVMFFIFieldFlagBitMaskDefaultFromFactory = 1 << 5,
873+
/*!
874+
* \brief The field is excluded from repr output.
875+
*
876+
* When set, the field will not appear in the generic reflection-based repr.
877+
* By default this flag is off (meaning the field is included in repr).
878+
*/
879+
kTVMFFIFieldFlagBitMaskReprOff = 1 << 6,
873880
#ifdef __cplusplus
874881
};
875882
#else

include/tvm/ffi/reflection/registry.h

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -223,6 +223,33 @@ class AttachFieldFlag : public InfoTrait {
223223
int32_t flag_;
224224
};
225225

226+
/*!
227+
* \brief Trait that controls whether a field appears in repr output.
228+
*
229+
* By default, all fields appear in repr. Use `Repr(false)` to exclude a field.
230+
*/
231+
class Repr : public InfoTrait {
232+
public:
233+
/*!
234+
* \brief Constructor.
235+
* \param show Whether the field should appear in repr output.
236+
*/
237+
explicit Repr(bool show) : show_(show) {}
238+
239+
/*!
240+
* \brief Apply the repr flag to the field info.
241+
* \param info The field info.
242+
*/
243+
TVM_FFI_INLINE void Apply(TVMFFIFieldInfo* info) const {
244+
if (!show_) {
245+
info->flags |= kTVMFFIFieldFlagBitMaskReprOff;
246+
}
247+
}
248+
249+
private:
250+
bool show_;
251+
};
252+
226253
/*!
227254
* \brief Get the byte offset of a class member field.
228255
*
@@ -493,6 +520,7 @@ struct init {
493520
namespace type_attr {
494521
inline constexpr const char* kInit = "__ffi_init__";
495522
inline constexpr const char* kShallowCopy = "__ffi_shallow_copy__";
523+
inline constexpr const char* kRepr = "__ffi_repr__";
496524
} // namespace type_attr
497525

498526
/*!

python/tvm_ffi/_ffi_api.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,6 +82,7 @@ def ModuleImportModule(_0: Module, _1: Module, /) -> None: ...
8282
def ModuleInspectSource(_0: Module, _1: str, /) -> str: ...
8383
def ModuleLoadFromFile(_0: str, /) -> Module: ...
8484
def ModuleWriteToFile(_0: Module, _1: str, _2: str, /) -> None: ...
85+
def ReprPrint(_0: Any, /) -> str: ...
8586
def Shape(*args: Any) -> Any: ...
8687
def String(_0: str, /) -> str: ...
8788
def StructuralKey(_0: Any, /) -> Any: ...
@@ -144,6 +145,7 @@ def ToJSONGraphString(_0: Any, _1: Any, /) -> str: ...
144145
"ModuleInspectSource",
145146
"ModuleLoadFromFile",
146147
"ModuleWriteToFile",
148+
"ReprPrint",
147149
"Shape",
148150
"String",
149151
"StructuralEqual",

python/tvm_ffi/container.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -197,7 +197,7 @@ def __repr__(self) -> str:
197197
# exception safety handling for chandle=None
198198
if self.__chandle__() == 0:
199199
return type(self).__name__ + "(chandle=None)"
200-
return "[" + ", ".join([x.__repr__() for x in self]) + "]"
200+
return str(core.__object_repr__(self)) # ty: ignore[unresolved-attribute]
201201

202202
def __contains__(self, value: object) -> bool:
203203
"""Check if the array contains a value."""
@@ -343,7 +343,7 @@ def __repr__(self) -> str:
343343
"""Return a string representation of the list."""
344344
if self.__chandle__() == 0:
345345
return type(self).__name__ + "(chandle=None)"
346-
return "[" + ", ".join([x.__repr__() for x in self]) + "]"
346+
return str(core.__object_repr__(self)) # ty: ignore[unresolved-attribute]
347347

348348
def __contains__(self, value: object) -> bool:
349349
"""Check if the list contains a value."""
@@ -543,4 +543,4 @@ def __repr__(self) -> str:
543543
# exception safety handling for chandle=None
544544
if self.__chandle__() == 0:
545545
return type(self).__name__ + "(chandle=None)"
546-
return "{" + ", ".join([f"{k.__repr__()}: {v.__repr__()}" for k, v in self.items()]) + "}"
546+
return str(core.__object_repr__(self)) # ty: ignore[unresolved-attribute]

python/tvm_ffi/cython/object.pxi

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,22 @@ def _set_class_object(cls):
2626
_CLASS_OBJECT = cls
2727

2828

29+
_REPR_PRINT = None
30+
_REPR_PRINT_LOADED = False
31+
32+
2933
def __object_repr__(obj: "Object") -> str:
30-
"""Object repr function that can be overridden by assigning to it"""
34+
"""Object repr function using ffi.ReprPrint when available."""
35+
global _REPR_PRINT, _REPR_PRINT_LOADED
36+
if not _REPR_PRINT_LOADED:
37+
_REPR_PRINT_LOADED = True
38+
_REPR_PRINT = _get_global_func("ffi.ReprPrint", False)
39+
if _REPR_PRINT is not None:
40+
try:
41+
return str(_REPR_PRINT(obj))
42+
except Exception: # noqa: BLE001
43+
# Silently fall back: __repr__ must never raise.
44+
pass
3145
return type(obj).__name__ + "(" + str(obj.__ctypes_handle__().value) + ")"
3246

3347

python/tvm_ffi/dataclasses/_utils.py

Lines changed: 0 additions & 40 deletions
Original file line numberDiff line numberDiff line change
@@ -127,46 +127,6 @@ def _get_all_fields(type_info: TypeInfo) -> list[TypeField]:
127127
return fields
128128

129129

130-
def method_repr(type_cls: type, type_info: TypeInfo) -> Callable[..., str]:
131-
"""Generate a ``__repr__`` method for the dataclass.
132-
133-
The generated representation includes all fields with ``repr=True`` in
134-
the format ``ClassName(field1=value1, field2=value2, ...)``.
135-
"""
136-
# Step 0. Collect all fields from the type hierarchy
137-
fields = _get_all_fields(type_info)
138-
139-
# Step 1. Filter fields that should appear in repr
140-
repr_fields: list[str] = []
141-
for field in fields:
142-
assert field.name is not None
143-
assert field.dataclass_field is not None
144-
if field.dataclass_field.repr:
145-
repr_fields.append(field.name)
146-
147-
# Step 2. Generate the repr method
148-
if not repr_fields:
149-
# No fields to show, return a simple class name representation
150-
body_lines = [f"return f'{type_cls.__name__}()'"]
151-
else:
152-
# Build field representations
153-
fields_str = ", ".join(
154-
f"{field_name}={{self.{field_name}!r}}" for field_name in repr_fields
155-
)
156-
body_lines = [f"return f'{type_cls.__name__}({fields_str})'"]
157-
158-
source_lines = ["def __repr__(self) -> str:"]
159-
source_lines.extend(f" {line}" for line in body_lines)
160-
source = "\n".join(source_lines)
161-
162-
# Note: Code generation in this case is guaranteed to be safe,
163-
# because the generated code does not contain any untrusted input.
164-
namespace: dict[str, Any] = {}
165-
exec(source, {}, namespace)
166-
__repr__ = namespace["__repr__"]
167-
return __repr__
168-
169-
170130
def method_init(_type_cls: type, type_info: TypeInfo) -> Callable[..., None]:
171131
"""Generate an ``__init__`` that forwards to the FFI constructor.
172132

python/tvm_ffi/dataclasses/c_class.py

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@
4141

4242
@dataclass_transform(field_specifiers=(field,), kw_only_default=False)
4343
def c_class(
44-
type_key: str, init: bool = True, kw_only: bool = False, repr: bool = True
44+
type_key: str, init: bool = True, kw_only: bool = False
4545
) -> Callable[[Type[_InputClsType]], Type[_InputClsType]]: # noqa: UP006
4646
"""(Experimental) Create a dataclass-like proxy for a C++ class registered with TVM FFI.
4747
@@ -77,10 +77,6 @@ def c_class(
7777
``__init__``. Individual fields can override this by setting
7878
``kw_only=False`` in :func:`field`. Additionally, a ``KW_ONLY`` sentinel
7979
annotation can be used to mark all subsequent fields as keyword-only.
80-
repr
81-
If ``True`` and the Python class does not define ``__repr__``, a
82-
representation method is auto-generated that includes all fields with
83-
``repr=True``.
8480
8581
Returns
8682
-------
@@ -128,9 +124,8 @@ class MyClass:
128124
"""
129125

130126
def decorator(super_type_cls: Type[_InputClsType]) -> Type[_InputClsType]: # noqa: UP006
131-
nonlocal init, repr
127+
nonlocal init
132128
init = init and "__init__" not in super_type_cls.__dict__
133-
repr = repr and "__repr__" not in super_type_cls.__dict__
134129
# Step 1. Retrieve `type_info` from registry
135130
type_info: TypeInfo = _lookup_or_register_type_info_from_type_key(type_key)
136131
assert type_info.parent_type_info is not None
@@ -146,11 +141,10 @@ def decorator(super_type_cls: Type[_InputClsType]) -> Type[_InputClsType]: # no
146141
)
147142
# Step 3. Create the proxy class with the fields as properties
148143
fn_init = _utils.method_init(super_type_cls, type_info) if init else None
149-
fn_repr = _utils.method_repr(super_type_cls, type_info) if repr else None
150144
type_cls: Type[_InputClsType] = _utils.type_info_to_cls( # noqa: UP006
151145
type_info=type_info,
152146
cls=super_type_cls,
153-
methods={"__init__": fn_init, "__repr__": fn_repr},
147+
methods={"__init__": fn_init},
154148
)
155149
_set_type_cls(type_info, type_cls)
156150
# Step 4. Set up __copy__, __deepcopy__, __replace__

python/tvm_ffi/dataclasses/field.py

Lines changed: 2 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -47,22 +47,20 @@ class Field:
4747
way the decorator understands.
4848
"""
4949

50-
__slots__ = ("default_factory", "init", "kw_only", "name", "repr")
50+
__slots__ = ("default_factory", "init", "kw_only", "name")
5151

5252
def __init__(
5353
self,
5454
*,
5555
name: str | None = None,
5656
default_factory: Callable[[], _FieldValue] | _MISSING_TYPE = MISSING,
5757
init: bool = True,
58-
repr: bool = True,
5958
kw_only: bool | _MISSING_TYPE = MISSING,
6059
) -> None:
6160
"""Do not call directly; use :func:`field` instead."""
6261
self.name = name
6362
self.default_factory = default_factory
6463
self.init = init
65-
self.repr = repr
6664
self.kw_only = kw_only
6765

6866

@@ -71,7 +69,6 @@ def field(
7169
default: _FieldValue | _MISSING_TYPE = MISSING,
7270
default_factory: Callable[[], _FieldValue] | _MISSING_TYPE = MISSING,
7371
init: bool = True,
74-
repr: bool = True,
7572
kw_only: bool | _MISSING_TYPE = MISSING,
7673
) -> _FieldValue:
7774
"""(Experimental) Declare a dataclass-style field on a :func:`c_class` proxy.
@@ -94,9 +91,6 @@ def field(
9491
init
9592
If ``True`` the field is included in the generated ``__init__``.
9693
If ``False`` the field is omitted from input arguments of ``__init__``.
97-
repr
98-
If ``True`` the field is included in the generated ``__repr__``.
99-
If ``False`` the field is omitted from the ``__repr__`` output.
10094
kw_only
10195
If ``True``, the field is a keyword-only argument in ``__init__``.
10296
If ``MISSING``, inherits from the class-level ``kw_only`` setting or
@@ -158,13 +152,11 @@ class PyBase:
158152
raise ValueError("Cannot specify both `default` and `default_factory`")
159153
if not isinstance(init, bool):
160154
raise TypeError("`init` must be a bool")
161-
if not isinstance(repr, bool):
162-
raise TypeError("`repr` must be a bool")
163155
if kw_only is not MISSING and not isinstance(kw_only, bool):
164156
raise TypeError(f"`kw_only` must be a bool, got {type(kw_only).__name__!r}")
165157
if default is not MISSING:
166158
default_factory = _make_default_factory(default)
167-
ret = Field(default_factory=default_factory, init=init, repr=repr, kw_only=kw_only)
159+
ret = Field(default_factory=default_factory, init=init, kw_only=kw_only)
168160
return cast(_FieldValue, ret)
169161

170162

0 commit comments

Comments
 (0)