helper functions for creating targeting flags

python/sdss_semaphore/__init__.py

@@ -2,15 +2,18 @@
 
 import numpy as np
 import warnings
-from typing import Union, Tuple, Iterable, List, Tuple
+from typing import Union, Tuple, Iterable, List, Optional, Tuple
 
 
 class BaseFlags:
 
     """A base class for communicating with flags."""
 
-    def __init__(self, array: Union[np.ndarray, Iterable[Iterable[int]], Iterable[bytearray]]) -> None:
-        if isinstance(array, (list, tuple)) and isinstance(array[0], bytearray):
+    def __init__(self, array: Optional[Union[np.ndarray, Iterable[Iterable[int]], Iterable[bytearray]]] = None) -> None:
+        if array is None:
+            # Assume single object flag.
+            self.array = np.zeros((1, 0), dtype=self.dtype)
+        elif isinstance(array, (list, tuple)) and isinstance(array[0], bytearray):
             # TODO: If the self.dtype is not uint8, then we might need to compute these initial offsets ourselves,
             #       because I think bytearray is natively uint8
             if self.dtype != np.uint8:

python/sdss_semaphore/targeting.py

@@ -1,5 +1,7 @@
 import numpy as np
 from typing import Tuple
+from functools import cached_property
+from typing import Optional
 from sdss_semaphore import BaseFlags
 
 
@@ -103,6 +105,27 @@ def count(self, skip_empty: bool = False) -> dict:
             skip_empty=skip_empty
         )
 
+    def set_bit_by_carton_pk(self, index: int, carton_pk: int):
+        """
+        Set the bit for the carton with the given primary key.
+
+        :param index:
+            The index of the item to set.
+
+        :param carton_pk:
+            The carton primary key.        
+        """
+        bit = self.get_bit_position_from_carton_pk[carton_pk]
+        return self.set_bit(index, bit)
+
+    @cached_property
+    def get_bit_position_from_carton_pk(self):
+        """
+        Return a dictionary with carton primary keys as keys, and bit positions as values.
+        
+        This is a helper method for efficiency creating large `TargetingFlags` objects.
+        """
+        return { attrs["carton_pk"]: bit for bit, attrs in self.mapping.items() }
 
 class TargetingFlags(BaseTargetingFlags):