added feature, to pass an index as mask to flags, because it seems a common workflow

saqc/flagger/flags.py

@@ -178,8 +178,15 @@ class Flags:
                 raise KeyError("a single 'column' or a tuple of 'mask, column' must be passt")
             mask, key = key
 
-            # raises (correct) KeyError
             tmp = pd.Series(UNTOUCHED, index=self._data[key].index, dtype=float)
+
+            # make a mask from an index, because it seems
+            # that passing an index is a very common workflow
+            if isinstance(mask, pd.Index):
+                mask = pd.Series(True, index=mask, dtype=bool)
+                mask = mask.reindex(tmp.index, fill_value=False)
+
+            # raises (correct) KeyError
             try:
                 tmp[mask] = value
             except Exception:

test/flagger/test_flags.py

@@ -188,6 +188,34 @@ def test_set_flags_with_mask(data: np.array):
                 flags[mask, c] = vector
 
 
+@pytest.mark.parametrize('data', data)
+def test_set_flags_with_index(data: np.array):
+    flags = Flags(data)
+
+    for c in flags.columns:
+        var = flags[c]
+        mask = var == UNFLAGGED
+        index = mask[mask].index
+
+        scalar = 222.
+        flags[index, c] = scalar
+        assert all(flags[c].loc[mask] == 222.)
+        assert all(flags[c].loc[~mask] != 222.)
+
+        vector = var.copy()
+        vector[:] = 333.
+        flags[index, c] = vector
+        assert all(flags[c].loc[mask] == 333.)
+        assert all(flags[c].loc[~mask] != 333.)
+
+        # works with any that pandas eat, eg with numpy
+        vector[:] = 444.
+        vector = vector.to_numpy()
+        flags[index, c] = vector
+        assert all(flags[c].loc[mask] == 444.)
+        assert all(flags[c].loc[~mask] != 444.)
+
+
 def test_cache():
     arr = np.array([
         [0, 0, 0, 0],