From 03e69596120bc6a42c53ca457252ecf70ffdbd3a Mon Sep 17 00:00:00 2001
From: William Jamieson <william@jamieson.phd>
Date: Tue, 12 Nov 2024 11:24:48 -0500
Subject: [PATCH] Bugfix for enums in python 3.10

---
 src/roman_datamodels/dqflags.py | 18 ++++++++++++++----
 tests/test_dqflags.py           | 11 +++++++++--
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/src/roman_datamodels/dqflags.py b/src/roman_datamodels/dqflags.py
index c6331be7..1b5fb266 100644
--- a/src/roman_datamodels/dqflags.py
+++ b/src/roman_datamodels/dqflags.py
@@ -19,14 +19,24 @@
 the formula `2**bit_number` where `bit_number` is the 0-index bit of interest.
 """
 
-from enum import Enum, unique
+import sys
+from enum import unique
 
-import numpy as np
+# Something with pickling of multiclassed enums was changed in 3.11 + allowing
+# us to directly us `np.uint32` as the enum object rather than a python `int`.
+if sys.version_info < (3, 11):
+    from enum import IntEnum
+else:
+    from enum import Enum
+
+    import numpy as np
+
+    class IntEnum(np.uint32, Enum): ...
 
 
 # fmt: off
 @unique
-class pixel(np.uint32, Enum):
+class pixel(IntEnum):
     """Pixel-specific data quality flags"""
 
     GOOD             = 0      # No bits set, all is good
@@ -64,7 +74,7 @@ class pixel(np.uint32, Enum):
 
 
 @unique
-class group(np.uint32, Enum):
+class group(IntEnum):
     """Group-specific data quality flags
         Once groups are combined, these flags are equivalent to the pixel-specific flags.
     """
diff --git a/tests/test_dqflags.py b/tests/test_dqflags.py
index 27694a0a..02bab7eb 100644
--- a/tests/test_dqflags.py
+++ b/tests/test_dqflags.py
@@ -1,3 +1,4 @@
+import sys
 from math import log10
 
 import numpy as np
@@ -31,7 +32,10 @@ def test_pixel_flags(flag):
     assert isinstance(flag, dqflags.pixel)
 
     # Test that the pixel flags are ints
-    assert isinstance(flag, np.uint32)
+    if sys.version_info < (3, 11):
+        assert isinstance(flag, int)
+    else:
+        assert isinstance(flag, np.uint32)
 
     # Test that the pixel flags are dict accessible
     assert dqflags.pixel[flag.name] is flag
@@ -79,7 +83,10 @@ def test_group_flags(flag):
     assert isinstance(flag, dqflags.group)
 
     # Test that the group flags are ints
-    assert isinstance(flag, np.uint32)
+    if sys.version_info < (3, 11):
+        assert isinstance(flag, int)
+    else:
+        assert isinstance(flag, np.uint32)
 
     # Test that the group flags are dict accessible
     assert dqflags.group[flag.name] is flag