Add AppendRemoveMatchRule class

Toilal · Oct 24, 2015 · 6a95ad6 · 6a95ad6
1 parent af983e2
commit 6a95ad6
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 1 deletion.
diff --git a/rebulk/__init__.py b/rebulk/__init__.py
@@ -6,5 +6,5 @@
 
 from .rebulk import Rebulk
 from .match import Match
-from .rules import Rule, AppendMatchRule, RemoveMatchRule
+from .rules import Rule, AppendMatchRule, RemoveMatchRule, AppendRemoveMatchRule
 from .pattern import REGEX_AVAILABLE
diff --git a/rebulk/rules.py b/rebulk/rules.py
@@ -95,6 +95,26 @@ def then(self, matches, when_response, context):
             matches.append(when_response)
 
 
+class AppendRemoveMatchRule(Rule):  # pylint: disable=abstract-method
+    """
+    Append matches returned by then[0] and remove matches returned by then[1]
+    """
+    def then(self, matches, when_response, context):
+        to_append, to_remove = when_response
+        if is_iterable(to_append):
+            for match in to_append:
+                matches.append(match)
+        else:
+            matches.append(to_append)
+        if is_iterable(to_remove):
+            for match in to_remove:
+                matches.remove(match)
+        else:
+            matches.append(to_remove)
+
+
+
+
 class Rules(list):
     """
     list of rules ready to execute.