This is an automated email from the ASF dual-hosted git repository.
junrushao 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 fc6cf8a fix(enum): add class iteration and len (#558)
fc6cf8a is described below
commit fc6cf8a0596b6cc1a7f3b453d18bab3c60310622
Author: Junru Shao <[email protected]>
AuthorDate: Mon Apr 20 12:29:20 2026 -0700
fix(enum): add class iteration and len (#558)
---
python/tvm_ffi/dataclasses/enum.py | 17 ++++++++++++++++-
tests/python/test_dataclass_enum.py | 10 ++++++++++
2 files changed, 26 insertions(+), 1 deletion(-)
diff --git a/python/tvm_ffi/dataclasses/enum.py
b/python/tvm_ffi/dataclasses/enum.py
index cc130ea..610c6c3 100644
--- a/python/tvm_ffi/dataclasses/enum.py
+++ b/python/tvm_ffi/dataclasses/enum.py
@@ -207,6 +207,21 @@ class _ClassProperty:
# ---------------------------------------------------------------------------
+class _EnumMeta(type(Object)):
+ """Metaclass for :class:`Enum` that wires class-level iteration and length.
+
+ Kept as a dedicated subclass of ``type(Object)`` so ``iter(MyEnum)`` and
+ ``len(MyEnum)`` only resolve on :class:`Enum` subclasses, rather than
+ mutating the shared ``Object`` metaclass globally.
+ """
+
+ def __iter__(cls) -> Iterator[Any]:
+ return iter(cls.by_value) # ty: ignore[unresolved-attribute]
+
+ def __len__(cls) -> int:
+ return len(cls.by_value) # ty: ignore[unresolved-attribute]
+
+
@dataclass_transform(
eq_default=False,
order_default=False,
@@ -214,7 +229,7 @@ class _ClassProperty:
field_specifiers=(Field, field, entry, auto),
)
@c_class("ffi.Enum", init=False)
-class Enum(Object):
+class Enum(Object, metaclass=_EnumMeta):
"""A named-singleton registry with cross-language identity.
Subclasses declare variants: frozen, named, ordinal-indexed
diff --git a/tests/python/test_dataclass_enum.py
b/tests/python/test_dataclass_enum.py
index af16601..d5a7925 100644
--- a/tests/python/test_dataclass_enum.py
+++ b/tests/python/test_dataclass_enum.py
@@ -118,6 +118,16 @@ def test_entries_iteration_order() -> None:
assert values[2].same_as(Ordered.c)
+def test_class_iteration_and_len() -> None:
+ class Ordered(Enum, type_key=_unique_key("OrderedIter")):
+ a = auto()
+ b = auto()
+ c = auto()
+
+ assert list(Ordered) == [Ordered.a, Ordered.b, Ordered.c]
+ assert len(Ordered) == 3
+
+
def test_frozen_variants() -> None:
class Frozen(Enum, type_key=_unique_key("Frozen")):
flag: bool
