address PR feedback to not use enums and instead use strings as a more pythonic version

rparolin · rparolin · commit 4c0c375a2094 · 2026-02-23T14:16:21.000-08:00
diff --git a/cuda_core/cuda/core/__init__.py b/cuda_core/cuda/core/__init__.py
@@ -37,7 +37,7 @@
     GraphCompleteOptions,
     GraphDebugPrintOptions,
 )
-from cuda.core._graphics import GraphicsRegisterFlags, GraphicsResource  # noqa: E402
+from cuda.core._graphics import GraphicsResource  # noqa: E402
 from cuda.core._launch_config import LaunchConfig  # noqa: E402
 from cuda.core._launcher import launch  # noqa: E402
 from cuda.core._layout import _StridedLayout  # noqa: E402
diff --git a/cuda_core/cuda/core/_graphics.pyx b/cuda_core/cuda/core/_graphics.pyx
@@ -13,29 +13,34 @@ from cuda.core._resource_handles cimport (
 from cuda.core._stream cimport Stream, Stream_accept
 from cuda.core._utils.cuda_utils cimport HANDLE_RETURN
 
-from enum import IntEnum
-
 from cuda.core._memory import Buffer
 
-__all__ = ['GraphicsResource', 'GraphicsRegisterFlags']
-
-
-class GraphicsRegisterFlags(IntEnum):
-    """Flags for registering a graphics resource with CUDA.
-
-    These flags specify the intended usage when registering a graphics
-    resource (e.g., an OpenGL buffer) for CUDA access.
-    """
-    NONE = cydriver.CU_GRAPHICS_REGISTER_FLAGS_NONE
-    """No hints about how this resource will be used. CUDA may read and write."""
-    READ_ONLY = cydriver.CU_GRAPHICS_REGISTER_FLAGS_READ_ONLY
-    """CUDA will not write to this resource."""
-    WRITE_DISCARD = cydriver.CU_GRAPHICS_REGISTER_FLAGS_WRITE_DISCARD
-    """CUDA will not read from this resource and will write over the entire contents."""
-    SURFACE_LOAD_STORE = cydriver.CU_GRAPHICS_REGISTER_FLAGS_SURFACE_LDST
-    """CUDA will bind this resource to a surface reference."""
-    TEXTURE_GATHER = cydriver.CU_GRAPHICS_REGISTER_FLAGS_TEXTURE_GATHER
-    """CUDA will perform texture gather operations on this resource."""
+__all__ = ['GraphicsResource']
+
+_REGISTER_FLAGS = {
+    "none": cydriver.CU_GRAPHICS_REGISTER_FLAGS_NONE,
+    "read_only": cydriver.CU_GRAPHICS_REGISTER_FLAGS_READ_ONLY,
+    "write_discard": cydriver.CU_GRAPHICS_REGISTER_FLAGS_WRITE_DISCARD,
+    "surface_load_store": cydriver.CU_GRAPHICS_REGISTER_FLAGS_SURFACE_LDST,
+    "texture_gather": cydriver.CU_GRAPHICS_REGISTER_FLAGS_TEXTURE_GATHER,
+}
+
+
+def _parse_register_flags(flags):
+    if flags is None:
+        return 0
+    if isinstance(flags, str):
+        flags = (flags,)
+    result = 0
+    for f in flags:
+        try:
+            result |= _REGISTER_FLAGS[f]
+        except KeyError:
+            raise ValueError(
+                f"Unknown register flag {f!r}. "
+                f"Valid flags: {', '.join(sorted(_REGISTER_FLAGS))}"
+            ) from None
+    return result
 
 
 class _MappedBufferContext:
@@ -114,16 +119,20 @@ cdef class GraphicsResource:
         )
 
     @classmethod
-    def from_gl_buffer(cls, int gl_buffer, *, int flags=0) -> GraphicsResource:
+    def from_gl_buffer(cls, int gl_buffer, *, flags=None) -> GraphicsResource:
         """Register an OpenGL buffer object for CUDA access.
 
         Parameters
         ----------
         gl_buffer : int
             The OpenGL buffer name (``GLuint``) to register.
-        flags : int, optional
-            Registration flags from :class:`GraphicsRegisterFlags`.
-            Defaults to :attr:`GraphicsRegisterFlags.NONE`.
+        flags : str or sequence of str, optional
+            Registration flags specifying intended usage. Accepted values:
+            ``"none"``, ``"read_only"``, ``"write_discard"``,
+            ``"surface_load_store"``, ``"texture_gather"``.
+            Multiple flags can be combined by passing a sequence
+            (e.g., ``("surface_load_store", "read_only")``).
+            Defaults to ``None`` (no flags).
 
         Returns
         -------
@@ -135,21 +144,24 @@ cdef class GraphicsResource:
         CUDAError
             If the registration fails (e.g., no current GL context, invalid
             buffer name, or operating system error).
+        ValueError
+            If an unknown flag string is provided.
         """
         cdef GraphicsResource self = GraphicsResource.__new__(cls)
         cdef cydriver.CUgraphicsResource resource
         cdef cydriver.GLuint cy_buffer = <cydriver.GLuint>gl_buffer
+        cdef unsigned int cy_flags = _parse_register_flags(flags)
         with nogil:
             HANDLE_RETURN(
-                cydriver.cuGraphicsGLRegisterBuffer(&resource, cy_buffer, <unsigned int>flags)
+                cydriver.cuGraphicsGLRegisterBuffer(&resource, cy_buffer, cy_flags)
             )
         self._handle = create_graphics_resource_handle(resource)
         self._mapped = False
         return self
 
     @classmethod
     def from_gl_image(
-        cls, int image, int target, *, int flags=0
+        cls, int image, int target, *, flags=None
     ) -> GraphicsResource:
         """Register an OpenGL texture or renderbuffer for CUDA access.
 
@@ -159,9 +171,13 @@ cdef class GraphicsResource:
             The OpenGL texture or renderbuffer name (``GLuint``) to register.
         target : int
             The OpenGL target type (e.g., ``GL_TEXTURE_2D``).
-        flags : int, optional
-            Registration flags from :class:`GraphicsRegisterFlags`.
-            Defaults to :attr:`GraphicsRegisterFlags.NONE`.
+        flags : str or sequence of str, optional
+            Registration flags specifying intended usage. Accepted values:
+            ``"none"``, ``"read_only"``, ``"write_discard"``,
+            ``"surface_load_store"``, ``"texture_gather"``.
+            Multiple flags can be combined by passing a sequence
+            (e.g., ``("surface_load_store", "read_only")``).
+            Defaults to ``None`` (no flags).
 
         Returns
         -------
@@ -172,14 +188,17 @@ cdef class GraphicsResource:
         ------
         CUDAError
             If the registration fails.
+        ValueError
+            If an unknown flag string is provided.
         """
         cdef GraphicsResource self = GraphicsResource.__new__(cls)
         cdef cydriver.CUgraphicsResource resource
         cdef cydriver.GLuint cy_image = <cydriver.GLuint>image
         cdef cydriver.GLenum cy_target = <cydriver.GLenum>target
+        cdef unsigned int cy_flags = _parse_register_flags(flags)
         with nogil:
             HANDLE_RETURN(
-                cydriver.cuGraphicsGLRegisterImage(&resource, cy_image, cy_target, <unsigned int>flags)
+                cydriver.cuGraphicsGLRegisterImage(&resource, cy_image, cy_target, cy_flags)
             )
         self._handle = create_graphics_resource_handle(resource)
         self._mapped = False
diff --git a/cuda_core/examples/gl_interop_plasma.py b/cuda_core/examples/gl_interop_plasma.py
@@ -64,7 +64,6 @@
 import numpy as np
 from cuda.core import (
     Device,
-    GraphicsRegisterFlags,
     GraphicsResource,
     LaunchConfig,
     Program,
@@ -296,7 +295,7 @@ def main():
     #     THIS IS THE KEY LINE.  GraphicsResource.from_gl_buffer() tells the
     #     CUDA driver "I want to access this OpenGL buffer from CUDA kernels."
     #     WRITE_DISCARD means CUDA will overwrite the entire buffer each frame.
-    resource = GraphicsResource.from_gl_buffer(pbo_id, flags=GraphicsRegisterFlags.WRITE_DISCARD)
+    resource = GraphicsResource.from_gl_buffer(pbo_id, flags="write_discard")
 
     # --- Step 6: Render loop ---
     start_time = time.monotonic()
diff --git a/cuda_core/tests/test_graphics.py b/cuda_core/tests/test_graphics.py
@@ -14,7 +14,6 @@
 from cuda.core import (
     Buffer,
     Device,
-    GraphicsRegisterFlags,
     GraphicsResource,
     StridedMemoryView,
 )
@@ -140,17 +139,33 @@ def _gl_context_and_texture(width=16, height=16):
 
 
 # ---------------------------------------------------------------------------
-# GraphicsRegisterFlags tests
+# Register flags parsing tests
 # ---------------------------------------------------------------------------
 
 
-class TestGraphicsRegisterFlags:
-    def test_enum_values(self):
-        assert int(GraphicsRegisterFlags.NONE) == 0
-        assert int(GraphicsRegisterFlags.READ_ONLY) == 1
-        assert int(GraphicsRegisterFlags.WRITE_DISCARD) == 2
-        assert int(GraphicsRegisterFlags.SURFACE_LOAD_STORE) == 4
-        assert int(GraphicsRegisterFlags.TEXTURE_GATHER) == 8
+class TestRegisterFlags:
+    def test_parse_none(self):
+        from cuda.core._graphics import _parse_register_flags
+
+        assert _parse_register_flags(None) == 0
+
+    def test_parse_single_string(self):
+        from cuda.core._graphics import _parse_register_flags
+
+        assert _parse_register_flags("read_only") == 1
+        assert _parse_register_flags("write_discard") == 2
+
+    def test_parse_combined_flags(self):
+        from cuda.core._graphics import _parse_register_flags
+
+        result = _parse_register_flags(("surface_load_store", "read_only"))
+        assert result == 4 | 1
+
+    def test_parse_invalid_raises(self):
+        from cuda.core._graphics import _parse_register_flags
+
+        with pytest.raises(ValueError, match="Unknown register flag"):
+            _parse_register_flags("bogus")
 
 
 # ---------------------------------------------------------------------------
@@ -179,7 +194,7 @@ def test_register_default_flags(self):
 
     def test_register_write_discard(self):
         with _gl_context_and_buffer() as (gl_buf, nbytes):
-            resource = GraphicsResource.from_gl_buffer(gl_buf, flags=GraphicsRegisterFlags.WRITE_DISCARD)
+            resource = GraphicsResource.from_gl_buffer(gl_buf, flags="write_discard")
             assert resource.handle != 0
             resource.close()
 
@@ -212,7 +227,7 @@ def test_register_image(self):
 class TestMapUnmap:
     def test_map_returns_buffer(self):
         with _gl_context_and_buffer(nbytes=4096) as (gl_buf, nbytes):
-            resource = GraphicsResource.from_gl_buffer(gl_buf, flags=GraphicsRegisterFlags.WRITE_DISCARD)
+            resource = GraphicsResource.from_gl_buffer(gl_buf, flags="write_discard")
             mapped = resource.map()
             assert resource.is_mapped
             # mapped is a _MappedBufferContext; its .handle and .size delegate to Buffer
@@ -224,7 +239,7 @@ def test_map_returns_buffer(self):
 
     def test_context_manager_unmaps(self):
         with _gl_context_and_buffer(nbytes=4096) as (gl_buf, nbytes):
-            resource = GraphicsResource.from_gl_buffer(gl_buf, flags=GraphicsRegisterFlags.WRITE_DISCARD)
+            resource = GraphicsResource.from_gl_buffer(gl_buf, flags="write_discard")
             with resource.map() as buf:
                 assert isinstance(buf, Buffer)
                 assert resource.is_mapped
@@ -234,7 +249,7 @@ def test_context_manager_unmaps(self):
 
     def test_context_manager_unmaps_on_exception(self):
         with _gl_context_and_buffer(nbytes=4096) as (gl_buf, nbytes):
-            resource = GraphicsResource.from_gl_buffer(gl_buf, flags=GraphicsRegisterFlags.WRITE_DISCARD)
+            resource = GraphicsResource.from_gl_buffer(gl_buf, flags="write_discard")
             with pytest.raises(ValueError, match="test error"), resource.map() as _buf:
                 assert resource.is_mapped
                 raise ValueError("test error")
@@ -246,7 +261,7 @@ def test_strided_memory_view_from_mapped_buffer(self):
         """End-to-end: register, map, create StridedMemoryView."""
         nbytes = 256 * 4  # 256 float32 elements
         with _gl_context_and_buffer(nbytes=nbytes) as (gl_buf, _):
-            resource = GraphicsResource.from_gl_buffer(gl_buf, flags=GraphicsRegisterFlags.WRITE_DISCARD)
+            resource = GraphicsResource.from_gl_buffer(gl_buf, flags="write_discard")
             with resource.map() as buf:
                 view = StridedMemoryView.from_buffer(buf, shape=(256,), dtype=np.float32)
                 assert view.ptr == int(buf.handle)
@@ -259,7 +274,7 @@ def test_map_with_stream(self):
             dev = Device(0)
             dev.set_current()
             stream = dev.create_stream()
-            resource = GraphicsResource.from_gl_buffer(gl_buf, flags=GraphicsRegisterFlags.WRITE_DISCARD)
+            resource = GraphicsResource.from_gl_buffer(gl_buf, flags="write_discard")
             with resource.map(stream=stream) as buf:
                 assert buf.size > 0
             resource.close()
@@ -304,7 +319,7 @@ def test_unmap_after_close_raises(self):
     def test_close_while_mapped(self):
         """close() should unmap before unregistering."""
         with _gl_context_and_buffer() as (gl_buf, nbytes):
-            resource = GraphicsResource.from_gl_buffer(gl_buf, flags=GraphicsRegisterFlags.WRITE_DISCARD)
+            resource = GraphicsResource.from_gl_buffer(gl_buf, flags="write_discard")
             resource.map()
             assert resource.is_mapped
             resource.close()  # Should unmap + unregister without error

Original file line number	Diff line number	Diff line change
`@@ -37,7 +37,7 @@`
`37`	`37`	`GraphCompleteOptions,`
`38`	`38`	`GraphDebugPrintOptions,`
`39`	`39`	`)`
`40`		`-from cuda.core._graphics import GraphicsRegisterFlags, GraphicsResource # noqa: E402`
	`40`	`+from cuda.core._graphics import GraphicsResource # noqa: E402`
`41`	`41`	`from cuda.core._launch_config import LaunchConfig # noqa: E402`
`42`	`42`	`from cuda.core._launcher import launch # noqa: E402`
`43`	`43`	`from cuda.core._layout import _StridedLayout # noqa: E402`