update code to have different classes

imap_processing/quality_flags.py

@@ -4,12 +4,26 @@
 
 
 class CommonFlags(IntFlag):
+    """Common quality flags."""
+
     NONE = 0x0
     INF = 2**0  # bit 0, Infinite value
+    NEG = 2**1  # bit 1, Negative value
+
+
+class ENAFlags(IntFlag):
+    """Common ENA flags."""
+
+    BADSPIN = 2**2  # bit 2, Bad spin
 
 
-class QualityFlags(IntFlag):
-    """Quality flags."""
+class ModifyFlags(IntFlag):
+    """
+    Quality flags modified.
+
+    This can be deleted once Python 3.11 is the
+    last support version of Python.
+    """
 
     @property
     def name(self) -> str:
@@ -24,54 +38,34 @@ def name(self) -> str:
         if self._name_ is not None:
             return self._name_
 
-        members = [member for member in QualityFlags if member & self == member]
-        return "|".join(
-            str(m).split(".", 1)[-1] for m in members if m != CommonFlags.NONE
-        )
+        members = [member for member in ModifyFlags if member & self == member]
+        return "|".join(str(m).split(".", 1)[-1] for m in members if m != 0x0)
 
 
-class NewFlag1(QualityFlags):
-    """New flag."""
+class ImapUltraFlags(ModifyFlags):
+    """IMAP Ultra flags."""
 
     NONE = CommonFlags.NONE
-    INF = CommonFlags.INF
-    FLAG2 = 2**1  # bit 1
+    INF = CommonFlags.INF  # bit 0
+    NEG = CommonFlags.NEG  # bit 1
+    BADSPIN = ENAFlags.BADSPIN  # bit 2
+    FLAG1 = 2**3  # bit 2
 
 
-class NewFlag2(QualityFlags):
-    """New flag."""
+class ImapLoFlags(ModifyFlags):
+    """IMAP Lo flags."""
 
     NONE = CommonFlags.NONE
-    INF = CommonFlags.INF
-    FLAG3 = 2**1  # bit 1
-
-
-# Assuming the previous code has been defined as shown
+    INF = CommonFlags.INF  # bit 0
+    NEG = CommonFlags.NEG  # bit 1
+    BADSPIN = ENAFlags.BADSPIN  # bit 2
+    FLAG2 = 2**3  # bit 2
 
-# Example usage of NewFlag1
-flag1 = NewFlag1.NONE
-print(f"NewFlag1.NONE: {flag1} ({flag1.name})")
 
-flag2 = NewFlag1.INF
-print(f"NewFlag1.INF: {flag2} ({flag2.name})")
+class HitFlags(ModifyFlags):
+    """Hit flags."""
 
-flag3 = NewFlag1.FLAG2
-print(f"NewFlag1.FLAG2: {flag3} ({flag3.name})")
-
-# Example of combining flags in NewFlag1
-combined_flags = NewFlag1.INF | NewFlag1.FLAG2
-print(f"Combined NewFlag1: {combined_flags} ({combined_flags.name})")
-
-# Example usage of NewFlag2
-flag4 = NewFlag2.NONE
-print(f"NewFlag2.NONE: {flag4} ({flag4.name})")
-
-flag5 = NewFlag2.INF
-print(f"NewFlag2.INF: {flag5} ({flag5.name})")
-
-flag6 = NewFlag2.FLAG3
-print(f"NewFlag2.FLAG3: {flag6} ({flag6.name})")
-
-# Example of combining flags in NewFlag2
-combined_flags2 = NewFlag2.INF | NewFlag2.FLAG3
-print(f"Combined NewFlag2: {combined_flags2} ({combined_flags2.name})")
+    NONE = CommonFlags.NONE
+    INF = CommonFlags.INF  # bit 0
+    NEG = CommonFlags.NEG  # bit 1
+    FLAG3 = 2**2  # bit 2

imap_processing/tests/test_quality_flags.py

@@ -1,44 +1,71 @@
-from imap_processing.quality_flags import QualityFlags
+"""Test bitwise flagging."""
+
+import numpy as np
+
+from imap_processing.quality_flags import HitFlags, ImapLoFlags, ImapUltraFlags
 
 
 def test_quality_flags():
-    """Test the QualityFlags bitwise operations."""
+    """Test the bitwise operations."""
 
     # Test individual flags
-    assert QualityFlags.NONE == 0x0
-    assert QualityFlags.INF == 2**0
-    assert QualityFlags.MISSING_TELEM == 2**1
-    assert QualityFlags.NEG == 2**2
-    assert QualityFlags.RES1 == 2**3
-    assert QualityFlags.BAD_SPIN == 2**4
-    assert QualityFlags.FOV == 2**5
-
-    flag = QualityFlags.INF | QualityFlags.RES1
-    assert flag & QualityFlags.INF
-    assert flag & QualityFlags.RES1
-    assert not flag & QualityFlags.MISSING_TELEM
-
-    assert QualityFlags.NONE.name == "NONE"
-    assert QualityFlags.INF.name == "INF"
-    combined_flags = QualityFlags.INF | QualityFlags.RES1
-    assert combined_flags.name == "INF|RES1"
-
-    combined_flags = QualityFlags.MISSING_TELEM | QualityFlags.FOV
-    assert combined_flags.name == "MISSING_TELEM|FOV"
-
-    combined_flags = (
-        QualityFlags.INF
-        | QualityFlags.MISSING_TELEM
-        | QualityFlags.NEG
-        | QualityFlags.RES1
+    assert HitFlags.NONE == 0x0
+    assert ImapUltraFlags.NONE == 0x0
+    assert ImapLoFlags.NONE == 0x0
+    assert HitFlags.INF == 2**0
+    assert ImapUltraFlags.INF == 2**0
+    assert ImapLoFlags.INF == 2**0
+    assert HitFlags.NEG == 2**1
+    assert ImapUltraFlags.NEG == 2**1
+    assert ImapLoFlags.NEG == 2**1
+    assert ImapUltraFlags.BADSPIN == 2**2
+    assert ImapLoFlags.BADSPIN == 2**2
+    assert ImapUltraFlags.FLAG1 == 2**3
+    assert ImapLoFlags.FLAG2 == 2**3
+    assert HitFlags.FLAG3 == 2**2
+
+    # Test combined flags for Ultra
+    flag = (
+        ImapUltraFlags.INF
+        | ImapUltraFlags.NEG
+        | ImapUltraFlags.BADSPIN
+        | ImapUltraFlags.FLAG1
     )
-    assert combined_flags.name == "INF|MISSING_TELEM|NEG|RES1"
+    assert flag & ImapUltraFlags.INF
+    assert flag & ImapUltraFlags.BADSPIN
+    assert flag & ImapUltraFlags.FLAG1
+    assert flag.name == "INF|NEG|BADSPIN|FLAG1"
+    assert flag.value == 15
 
-    combined_flags = QualityFlags.BAD_SPIN | QualityFlags.FOV | QualityFlags.INF
-    assert combined_flags.name == "INF|BAD_SPIN|FOV"
+    # Test combined flags for Lo
+    flag = ImapLoFlags.INF | ImapLoFlags.NEG | ImapLoFlags.BADSPIN | ImapLoFlags.FLAG2
+    assert flag & ImapLoFlags.INF
+    assert flag & ImapLoFlags.BADSPIN
+    assert flag & ImapLoFlags.FLAG2
+    assert flag.name == "INF|NEG|BADSPIN|FLAG2"
+    assert flag.value == 15
 
-    combined_flags = QualityFlags.FOV | QualityFlags.RES1 | QualityFlags.MISSING_TELEM
-    assert combined_flags.name == "MISSING_TELEM|RES1|FOV"
+    # Test combined flags for HIT
+    flag = HitFlags.INF | HitFlags.NEG | HitFlags.FLAG3
+    assert flag & HitFlags.INF
+    assert flag & HitFlags.FLAG3
+    assert flag.name == "INF|NEG|FLAG3"
+    assert flag.value == 7
+
+    # Test use-case for Ultra
+    data = np.array([-6, np.inf, 2, 3])
+    quality = np.array(
+        [
+            ImapUltraFlags.INF | ImapUltraFlags.NEG,
+            ImapUltraFlags.INF,
+            ImapUltraFlags.NONE,
+            ImapUltraFlags.NONE,
+        ]
+    )
+    # Select data without INF flags
+    non_inf_mask = (quality & ImapUltraFlags.INF.value) == 0
+    np.array_equal(data[non_inf_mask], np.array([-6, 2, 3]))
 
-    combined_flags = QualityFlags.BAD_SPIN | QualityFlags.FOV
-    assert combined_flags.name == "BAD_SPIN|FOV"
+    # Select data without NEG or INF flags
+    non_neg_mask = (quality & ImapUltraFlags.NEG.value) == 0
+    np.array_equal(data[non_inf_mask & non_neg_mask], np.array([2, 3]))