numba-cuda 0.19.1__py3-none-any.whl → 0.20.1__py3-none-any.whl

numba_cuda/numba/cuda/typing/enumdecl.py

@@ -0,0 +1,74 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: BSD-2-Clause
+
+"""
+Typing for enums.
+"""
+
+import operator
+from numba.core import types
+from numba.cuda.typing.templates import (
+    AbstractTemplate,
+    AttributeTemplate,
+    signature,
+    Registry,
+)
+
+registry = Registry()
+infer = registry.register
+infer_global = registry.register_global
+infer_getattr = registry.register_attr
+
+
+@infer_getattr
+class EnumAttribute(AttributeTemplate):
+    key = types.EnumMember
+
+    def resolve_value(self, ty):
+        return ty.dtype
+
+
+@infer_getattr
+class EnumClassAttribute(AttributeTemplate):
+    key = types.EnumClass
+
+    def generic_resolve(self, ty, attr):
+        """
+        Resolve attributes of an enum class as enum members.
+        """
+        if attr in ty.instance_class.__members__:
+            return ty.member_type
+
+
+@infer
+class EnumClassStaticGetItem(AbstractTemplate):
+    key = "static_getitem"
+
+    def generic(self, args, kws):
+        enum, idx = args
+        if (
+            isinstance(enum, types.EnumClass)
+            and idx in enum.instance_class.__members__
+        ):
+            return signature(enum.member_type, *args)
+
+
+class EnumCompare(AbstractTemplate):
+    def generic(self, args, kws):
+        [lhs, rhs] = args
+        if (
+            isinstance(lhs, types.EnumMember)
+            and isinstance(rhs, types.EnumMember)
+            and lhs == rhs
+        ):
+            return signature(types.boolean, lhs, rhs)
+
+
+@infer_global(operator.eq)
+class EnumEq(EnumCompare):
+    pass
+
+
+@infer_global(operator.ne)
+class EnumNe(EnumCompare):
+    pass