added row-select support on set flags

saqc/flagger/flags.py

@@ -6,7 +6,7 @@ import dios
 from saqc.common import *
 from saqc.flagger.history import History
 import pandas as pd
-from typing import Union, Dict, DefaultDict, Optional, Type
+from typing import Union, Dict, DefaultDict, Optional, Type, Tuple, Iterable
 
 _VAL = Union[pd.Series, History]
 DictLike = Union[
@@ -16,6 +16,13 @@ DictLike = Union[
     DefaultDict[str, _VAL],
 ]
 
+_Field = str
+SelectT = Union[
+    _Field,
+    Tuple[pd.Series, _Field]
+]
+ValueT = Union[pd.Series, Iterable, float]
+
 
 class _HistAccess:
 
@@ -163,18 +170,38 @@ class Flags:
 
         return self._cache[key].copy()
 
-    def __setitem__(self, key: str, value: pd.Series, force=False):
+    def __setitem__(self, key: SelectT, value: ValueT, force=False):
         # force-KW is internal available only
 
+        if isinstance(key, tuple):
+            if len(key) != 2:
+                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)
+            try:
+                tmp[mask] = value
+            except Exception:
+                raise ValueError('bad mask')
+            else:
+                value = tmp
+
+        # technically it would be possible to select a field and set
+        # the entire column to a scalar flag value (float), but it has
+        # a high potential, that this is not intended by the user.
+        if not isinstance(value, pd.Series):
+            raise ValueError("must pass value of type pd.Series")
+
         # if nothing happens no-one writes the history books
-        if isinstance(value, pd.Series) and len(value) == 0:
+        if len(value) == 0:
             return
 
         if key not in self._data:
             self._data[key] = History()
-        
+
         self._data[key].append(value, force=force)
-        self._cache.pop(key, None)     
+        self._cache.pop(key, None)
 
     def force(self, key: str, value: pd.Series) -> Flags:
         """
@@ -314,4 +341,3 @@ def initFlagsLike(reference: Union[pd.Series, DictLike, Flags], initial_value: f
 
 # for now we keep this name
 Flagger = Flags
-

test/flagger/test_flags.py

@@ -3,8 +3,8 @@ import dios
 import pytest
 import numpy as np
 import pandas as pd
-from pandas.api.types import is_bool_dtype
-from test.common import TESTFLAGGER, initData
+
+from saqc import BAD, UNFLAGGED
 from test.flagger.test_history import (
     History,
     is_equal as hist_equal,
@@ -156,8 +156,42 @@ def test_set_flags_and_force(data: np.array):
 
 
 @pytest.mark.parametrize('data', data)
-def test_force_flags(data: np.array):
-    pass
+def test_set_flags_with_mask(data: np.array):
+    flags = Flags(data)
+
+    for c in flags.columns:
+        var = flags[c]
+        mask = var == UNFLAGGED
+
+        scalar = 222.
+        flags[mask, c] = scalar
+        assert all(flags[c].loc[mask] == 222.)
+        assert all(flags[c].loc[~mask] != 222.)
+
+        # scalar without mask is not allowed, because
+        # it holds to much potential to set the whole
+        # column unintentionally.
+        with pytest.raises(ValueError):
+            flags[c] = 888.
+
+        vector = var.copy()
+        vector[:] = 333.
+        flags[mask, 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[mask, c] = vector
+        assert all(flags[c].loc[mask] == 444.)
+        assert all(flags[c].loc[~mask] != 444.)
+
+        # test length miss-match
+        if len(vector):
+            vector = vector[:-1]
+            with pytest.raises(ValueError):
+                flags[mask, c] = vector
 
 
 def test_cache():

test/flagger/test_history.py

@@ -3,8 +3,6 @@
 import pytest
 import numpy as np
 import pandas as pd
-from pandas.api.types import is_bool_dtype
-from test.common import TESTFLAGGER, initData
 from saqc.flagger.history import History
 
 # see #GH143 combined backtrack