fixed all

saqc/flagger/baseflagger.py

@@ -64,22 +64,24 @@ class BaseFlagger:
     def getFlags(self, flags: PandasLike) -> PandasLike:
         return flags
 
-    def setFlags(self, flags: pd.DataFrame, field, mask_or_indexer=None, flag=None, **kwargs) -> pd.DataFrame:
+    def setFlags(self, flags: pd.DataFrame, field, loc=None, iloc=None, flag=None, **kwargs) -> pd.DataFrame:
         if not isinstance(flags, pd.DataFrame):
             raise TypeError(f"flags must be of type pd.DataFrame, {type(flags)} was given")
         # prepare
-        flags = self._assureDtype(flags.copy(), field)
-        r = slice(None) if mask_or_indexer is None else mask_or_indexer
+        flags = self._assureDtype(flags, field)
         flag = self.BAD if flag is None else self._checkFlag(flag)
+        flags_loc, rows, col = self._getIndexer(flags, field, loc, iloc)
         # set
-        mask = flags.loc[r, field] < flag
+        mask = flags_loc[rows, col] < flag
         idx = mask[mask].index
         flags.loc[idx, field] = flag
         return self._assureDtype(flags, field)
 
-    def clearFlags(self, flags, field, mask_or_indexer=None, **kwargs):
-        moi = slice(None) if mask_or_indexer is None else mask_or_indexer
-        flags.loc[moi, field] = self.UNFLAGGED
+    def clearFlags(self, flags, field, loc=None, iloc=None, **kwargs):
+        if not isinstance(flags, pd.DataFrame):
+            raise TypeError(f"flags must be of type pd.DataFrame, {type(flags)} was given")
+        flags_loc, rows, col = self._getIndexer(flags, field, loc, iloc)
+        flags_loc[rows, col] = self.UNFLAGGED
         return self._assureDtype(flags, field)
 
     def _checkFlag(self, flag):
@@ -92,8 +94,19 @@ class BaseFlagger:
                 raise ValueError(f"Invalid flag '{flag}'. Possible choices are {list(self.flags.categories)[1:]}")
         return flag
 
+    def _getIndexer(self, flags, field, loc=None, iloc=None):
+        if loc is not None and iloc is not None:
+            raise ValueError("params `loc` and `iloc` are mutual exclusive")
+        elif loc is not None and iloc is None:
+            indexer, rows, col = flags.loc, loc, field
+        elif loc is None and iloc is not None:
+            indexer, rows, col = flags.iloc, iloc, flags.columns.get_loc(field)
+        elif loc is None and iloc is None:
+            indexer, rows, col = flags.loc, slice(None), field
+        return indexer, rows, col
+
     def _assureDtype(self, flags, field=None):
-        if field is None:
+        if field is None:  # we got a df
             flags = flags.astype(self.flags)
         elif not isinstance(flags[field].dtype, pd.Categorical):
             flags[field] = flags[field].astype(self.flags)

saqc/flagger/dmpflagger.py

@@ -50,15 +50,19 @@ class DmpFlagger(BaseFlagger):
         return super().isFlagged(flags, flag, comparator)
 
     def getFlags(self, flags: PandasLike) -> PandasLike:
-        if not isinstance(flags, pd.DataFrame):
+        if not isinstance(flags, pd.DataFrame):  # series
             super().getFlags(flags)
-        if isinstance(flags.columns, pd.MultiIndex):
+        elif isinstance(flags.columns, pd.MultiIndex):  # df, multiindex
             flags = flags.xs(FlagFields.FLAG, level=ColumnLevels.FLAGS, axis=1)
-        else:
+        else:  # df, simple index
+            # if we come here a dataframe with multiindex, was unpacked by a fieldname, like so: getFlags(df[field])
+            # so we can be sure to have qflag, qcause, qcomment as df-columns (otherwise something went wrong on the
+            # user side). As so we need to squeeze the result to a series, because the df wont have any column info,
+            # like to which field/variable the qflags belongs to.
             flags = flags[FlagFields.FLAG].squeeze()
         return flags
 
-    def setFlags(self, flags, field, mask_or_indexer=None, flag=None, comment='', cause='', **kwargs):
+    def setFlags(self, flags, field, loc=None, iloc=None, flag=None, comment='', cause='', **kwargs):
         if not isinstance(flags, pd.DataFrame):
             raise TypeError(f"flags must be of type pd.DataFrame, {type(flags)} was given")
         if not isinstance(flags.columns, pd.MultiIndex):
@@ -67,34 +71,31 @@ class DmpFlagger(BaseFlagger):
         ysoncomment = json.dumps(dict(comment=comment, commit=self.project_version, test=kwargs.get("func_name", "")))
         comment = '' if comment is None else ysoncomment
         flags = self._assureDtype(flags, field)
-        moi = slice(None) if mask_or_indexer is None else mask_or_indexer
         flag = self.BAD if flag is None else self._checkFlag(flag)
         # set
-        f = self.getFlags(flags[field])
+        indexer, rows, col = self._getIndexer(self.getFlags(flags), field, loc, iloc)
         if isinstance(flag, pd.Series):
-            flag = flag.loc[moi]
-        mask = f.loc[moi] < flag
+            i, r, _ = self._getIndexer(flag, field, loc, iloc)
+            flag = i[r]
+        mask = indexer[rows, col] < flag
         idx = mask[mask].index
         flags.loc[idx, field] = flag, cause, comment
         return self._assureDtype(flags, field)
 
-    def clearFlags(self, flags, field, mask_or_indexer=None, **kwargs):
+    def clearFlags(self, flags, field, loc=None, iloc=None, **kwargs):
         if not isinstance(flags, pd.DataFrame):
             raise TypeError(f"flags must be of type pd.DataFrame, {type(flags)} was given")
         if not isinstance(flags.columns, pd.MultiIndex):
             raise TypeError(f"flags.index must be a multiindex")
-        # prepare
-        flags = self._assureDtype(flags, field)
-        moi = slice(None) if mask_or_indexer is None else mask_or_indexer
-        # set
-        flags.loc[moi, field] = self.UNFLAGGED, '', ''
+        indexer, rows, col = self._getIndexer(flags, field, loc, iloc)
+        indexer[rows, col] = self.UNFLAGGED, '', ''
         return self._assureDtype(flags, field)
 
     def _assureDtype(self, flags, field=None):
         if field is None:
-            if isinstance(flags, pd.Series):
+            if isinstance(flags, pd.Series):  # we got a series
                 flags = super()._assureDtype(flags, None)
-            else:
+            else:  # we got a df with multiindex
                 flags = flags.astype({c: self.flags for c in flags.columns if FlagFields.FLAG in c})
         elif not isinstance(flags[(field, FlagFields.FLAG)].dtype, pd.Categorical):
             flags[(field, FlagFields.FLAG)] = flags[(field, FlagFields.FLAG)].astype(self.flags)

test/dsl/test_evaluator.py

@@ -48,7 +48,7 @@ def test_flagPropagation():
     data = initData()
     flagger = SimpleFlagger()
     flags = flagger.initFlags(data)
-    flags.iloc[::5] = flagger.setFlag(flags.iloc[::5])
+    flags = flagger.setFlags(flags, 'var2', iloc=slice(None, None, 5))
 
     var1, var2, *_ = data.columns
     var2_flags = flagger.isFlagged(flags[var2])
@@ -67,8 +67,8 @@ def test_isflagged():
     data = initData(cols=1)
     flagger = SimpleFlagger()
     flags = flagger.initFlags(data)
-    flags.iloc[::5] = flagger.setFlag(flags.iloc[::5], flagger.BAD)
-    flags.iloc[1::5] = flagger.setFlag(flags.iloc[1::5], flagger.GOOD)
+    flags = flagger.setFlags(flags, 'var1', iloc=slice(None, None, 5), flag=flagger.BAD)
+    flags = flagger.setFlags(flags, 'var1', iloc=slice(1, None, 5), flag=flagger.GOOD)
 
     tests = {
         "isflagged(this)" : flagger.isFlagged(flags, flagger.GOOD, ">"),

test/dsl/test_generic.py

@@ -43,12 +43,10 @@ def test_isflagged():
     flags = flagger.initFlags(data)
     var1, var2, *_ = data.columns
 
-    flags.iloc[::2, 0] = flagger.setFlag(flags.iloc[::2, 0])
+    flags = flagger.setFlags(flags, var1, iloc=slice(None, None, 2))
+    flags = flagger.setFlags(flags, var2, iloc=slice(None, None, 2))
 
-    idx = evalExpression("isflagged({:})".format(var1),
-                         flagger,
-                         data, flags,
-                         var2)
+    idx = evalExpression(f"isflagged({var1})", flagger, data, flags, var2)
 
     flagged = flagger.isFlagged(flags[var1])
     assert (flagged == idx).all
@@ -61,12 +59,9 @@ def test_isflaggedArgument():
     flags = flagger.initFlags(data)
     var1, var2, *_ = data.columns
 
-    flags.iloc[::2, 0] = flagger.setFlag(flags.iloc[::2, 0], 1)
+    flags = flagger.setFlags(flags, var1, iloc=slice(None, None, 2), flag=1)
 
-    idx = evalExpression("isflagged({:}, 1)".format(var1),
-                         flagger,
-                         data, flags,
-                         var2)
+    idx = evalExpression(f"isflagged({var1}, 1)", flagger, data, flags, var2)
 
     flagged = flagger.isFlagged(flags[var1], 1)
     assert (flagged == idx).all