reworked setFlag(s)

saqc/flagger/baseflagger.py

@@ -64,6 +64,17 @@ class BaseFlagger:
     def getFlags(self, flags: PandasLike):
         return flags
 
+    def setFlags(self, flags, field, mask_or_indexer=None, flag=None, **kwargs):
+        # prepare
+        flags = flags.copy()
+        r = slice(None) if mask_or_indexer is None else mask_or_indexer
+        flag = self.BAD if flag is None else self._checkFlag(flag)
+        # set
+        mask = flags.loc[r, field] < flag
+        idx = mask[mask].index
+        flags.loc[idx, field] = flag
+        return flags
+
     def setFlag(self, flags: pd.Series, flag: Optional[T] = None, **kwargs: Any) -> pd.Series:
         flags = self._checkFlagsType(flags)
         flag = self.BAD if flag is None else self._checkFlag(flag)

saqc/flagger/dmpflagger.py

@@ -54,9 +54,22 @@ class DmpFlagger(BaseFlagger):
         if isinstance(flags.columns, pd.MultiIndex):
             flags = flags.xs(FlagFields.FLAG, level=ColumnLevels.FLAGS, axis=1)
         else:
-            flags = flags[FlagFields.FLAG]
-        if isinstance(flags, pd.DataFrame):
-            flags = flags.to_frame()
+            flags = flags[FlagFields.FLAG].squeeze()
+        return flags
+
+    def setFlags(self, flags, field, mask_or_indexer=None, flag=None, comment='', cause='', **kwargs):
+        comment = json.dumps({
+            "comment": comment,
+            "commit": self.project_version,
+            "test": kwargs.get("func_name", "")})
+
+        flags = flags.copy()
+        r = slice(None) if mask_or_indexer is None else mask_or_indexer
+        flag = self.BAD if flag is None else self._checkFlag(flag)
+        f = self.getFlags(flags[field])
+        mask = f.loc[r] < flag
+        idx = mask[mask].index
+        flags.loc[idx, field] = flag, cause, comment
         return flags
 
     def setFlag(self, flags, flag=None, cause="", comment="", **kwargs):