Improve _EnumType and _Enum

Author: Luthaf

scripts/update-declarations.py

@@ -36,10 +36,7 @@ def __new__(metacls, name, bases, namespace):
             members = namespace["_members_"]
 
         namespace["_reverse_map_"] = {v: k for k, v in members.items()}
-        cls = type(ctypes.c_int32).__new__(metacls, name, bases, namespace)
-        for key, value in cls._members_.items():
-            globals()[key] = value
-        return cls
+        return type(ctypes.c_int32).__new__(metacls, name, bases, namespace)
 
     def __repr__(self):
         return f"<Enum {self.__name__}>"
@@ -55,8 +52,12 @@ def __repr__(self):
     def __eq__(self, other):
         if isinstance(other, int):
             return self.value == other
+        if type(self) is type(other):
+            return self.value == other.value
+        return NotImplemented
 
-        return type(self) is type(other) and self.value == other.value
+    def __hash__(self):
+        return hash(self.value)
 """
 
 

src/ctypes_dlpack/_c_api.py

@@ -20,10 +20,7 @@ def __new__(metacls, name, bases, namespace):
             members = namespace["_members_"]
 
         namespace["_reverse_map_"] = {v: k for k, v in members.items()}
-        cls = type(ctypes.c_int32).__new__(metacls, name, bases, namespace)
-        for key, value in cls._members_.items():
-            globals()[key] = value
-        return cls
+        return type(ctypes.c_int32).__new__(metacls, name, bases, namespace)
 
     def __repr__(self):
         return f"<Enum {self.__name__}>"
@@ -39,8 +36,12 @@ def __repr__(self):
     def __eq__(self, other):
         if isinstance(other, int):
             return self.value == other
+        if type(self) is type(other):
+            return self.value == other.value
+        return NotImplemented
 
-        return type(self) is type(other) and self.value == other.value
+    def __hash__(self):
+        return hash(self.value)
 
 
 class DLDeviceType(_Enum):

tests/test_enum.py

@@ -0,0 +1,39 @@
+from ctypes_dlpack import DLDataTypeCode, DLDeviceType
+
+
+def test_enum_eq():
+    assert DLDeviceType.kDLCPU == 1
+    assert 1 == DLDeviceType.kDLCPU
+    assert DLDeviceType.kDLCUDA == 2
+    assert DLDeviceType.kDLCPU != 2
+
+    assert DLDeviceType.kDLCPU == DLDeviceType.kDLCPU
+    assert DLDeviceType.kDLCPU != DLDeviceType.kDLCUDA
+
+    # Both have value 1, but different enum types — must not compare equal.
+    assert DLDeviceType.kDLCPU != DLDataTypeCode.kDLInt
+    assert DLDataTypeCode.kDLInt != DLDeviceType.kDLCPU
+
+
+def test_enum_hashable():
+    table = {DLDeviceType.kDLCPU: "cpu", DLDeviceType.kDLCUDA: "cuda"}
+    assert table[DLDeviceType.kDLCPU] == "cpu"
+    assert table[1] == "cpu"
+    assert table[DLDeviceType.kDLCUDA] == "cuda"
+
+    s = {DLDeviceType.kDLCPU, DLDeviceType.kDLCUDA, DLDeviceType.kDLCPU}
+    assert len(s) == 2
+
+
+def test_enum_repr():
+    # ``DLDeviceType.kDLCPU`` is a plain ``int`` (class-body assignment); the
+    # custom ``__repr__`` only applies to instances, as produced by ctypes
+    # when reading a structure field typed as ``DLDeviceType``.
+    assert repr(DLDeviceType(1)) == "DLDeviceType.kDLCPU"
+    assert repr(DLDataTypeCode(2)) == "DLDataTypeCode.kDLFloat"
+
+    unknown = DLDeviceType(99)
+    assert repr(unknown) == "DLDeviceType.99"
+
+    assert repr(DLDeviceType) == "<Enum DLDeviceType>"
+    assert repr(DLDataTypeCode) == "<Enum DLDataTypeCode>"