prepared flagger Base, Dmp

saqc/flagger/baseflagger.py

@@ -164,6 +164,9 @@ class BaseFlagger(ABC):
             force: bool = False,
             inplace=False,
             with_extra=False,
+            flag_after=None,
+            flag_before=None,
+            win_flag=None,
             **kwargs
     ) -> BaseFlaggerT:
         """Overwrite existing flags at loc.
@@ -192,16 +195,32 @@ class BaseFlagger(ABC):
         if with_extra and not isinstance(flag, pd.Series):
             raise ValueError("flags must be pd.Series if `with_extras=True`.")
 
+        trimmed = self.getFlags(field=field, loc=loc)
         if force:
-            row_indexer = slice(None) if loc is None else loc
+            mask = pd.Series(True, index=trimmed.index, dtype=bool)
         else:
-            # trim flags to loc, we always get a pd.Series returned
-            this = self.getFlags(field=field, loc=loc)
-            row_indexer = this < flag
+            mask = trimmed < flag
+
+        # set flags of the test
+        out._flags.aloc[mask, field] = flag
+
+        # calc and set window flags
+        if flag_after is not None or flag_before is not None:
+            win_mask = self._getWindowMask(field, mask, flag_after, flag_before)
+            out._flags.aloc[win_mask, field] = win_flag
 
-        out._flags.aloc[row_indexer, field] = flag
         return out
 
+    def _getWindowMask(self, field, mask, flag_after, flag_before):
+        win_mask = mask.reindex_like(self._flags[field]).fillna(False)
+        # does not overwrite the flags from the test again
+        skip = ~win_mask
+        if flag_after is not None:
+            win_mask |= pd.Series(False)
+        if flag_before is not None:
+            win_mask |= pd.Series(False)
+        return win_mask & skip
+
     def clearFlags(self, field: str, loc: LocT = None, inplace=False, **kwargs) -> BaseFlaggerT:
         assertScalar("field", field, optional=False)
         if "force" in kwargs:

saqc/flagger/dmpflagger.py

@@ -132,6 +132,9 @@ class DmpFlagger(CategoricalFlagger):
         force=False,
         inplace=False,
         with_extra=False,
+        flag_after=None,
+        flag_before=None,
+        win_flag=None,
         **kwargs
     ):
         assert "iloc" not in kwargs, "deprecated keyword, iloc"
@@ -153,15 +156,24 @@ class DmpFlagger(CategoricalFlagger):
                  "test": kwargs.get("func_name", "")}
             )
 
+        flags = self.getFlags(field=field, loc=loc)
         if force:
-            row_indexer = slice(None) if loc is None else loc
+            mask = pd.Series(True, index=flags.index, dtype=bool)
         else:
-            this = self.getFlags(field=field, loc=loc)
-            row_indexer = this < flag
+            mask = flags < flag
+
+        # set flags of the test
+        out._flags.aloc[mask, field] = flag
+        out._causes.aloc[mask, field] = cause
+        out._comments.aloc[mask, field] = comment
+
+        # calc and set window flags
+        if flag_after is not None or flag_before is not None:
+            win_mask = self._getWindowMask(field, mask, flag_after, flag_before)
+            out._flags.aloc[win_mask, field] = win_flag
+            out._causes.aloc[win_mask, field] = cause
+            out._comments.aloc[win_mask, field] = comment
 
-        out._flags.aloc[row_indexer, field] = flag
-        out._causes.aloc[row_indexer, field] = cause
-        out._comments.aloc[row_indexer, field] = comment
         return out
 
     def replaceField(self, field, flags, inplace=False, cause=None, comment=None, **kwargs):