Hold on to BATCH args, allow event hooks to manipulate batches

2 files changed, 16 insertions, 5 deletions

diff --git a/modules/line_handler/__init__.py b/modules/line_handler/__init__.py
index 70e53d43..73057b55 100644
--- a/modules/line_handler/__init__.py
+++ b/modules/line_handler/__init__.py
@@ -182,7 +182,10 @@ class Module(ModuleManager.BaseModule):
 
         if modifier == "+":
             batch_type = event["args"][1]
-            batch = utils.irc.IRCBatch(identifier, batch_type, event["tags"])
+            args = event["args"][2:]
+
+            batch = utils.irc.IRCBatch(identifier, batch_type, args,
+                event["tags"])
             event["server"].batches[identifier] = batch
 
             self.events.on("received.batch.start").call(batch=batch,
@@ -191,10 +194,17 @@ class Module(ModuleManager.BaseModule):
             batch = event["server"].batches[identifier]
             del event["server"].batches[identifier]
 
-            self.events.on("received.batch.end").call(batch=batch,
+            lines = batch.get_lines()
+
+            results = self.events.on("received.batch.end").call(batch=batch,
                 server=event["server"])
 
-            for line in batch.lines:
+            for result in results:
+                if not result == None:
+                    lines = result
+                    break
+
+            for line in lines:
                 self._handle(event["server"], line)
 
     # IRCv3 CHGHOST, a user's username and/or hostname has changed
diff --git a/src/utils/irc/__init__.py b/src/utils/irc/__init__.py
index 3e3cef1b..5a6fb69b 100644
--- a/src/utils/irc/__init__.py
+++ b/src/utils/irc/__init__.py
@@ -258,10 +258,11 @@ def parse_ctcp(s: str) -> typing.Optional[CTCPMessage]:
     return None
 
 class IRCBatch(object):
-    def __init__(self, identifier: str, batch_type: str, tags:
-            typing.Dict[str, str]={}):
+    def __init__(self, identifier: str, batch_type: str, args: typing.List[str],
+            tags: typing.Dict[str, str]={}):
         self.id = identifier
         self.type = batch_type
+        self.args = args
         self.tags = tags
         self._lines = [] # type: typing.List[IRCLine.ParsedLine]
     def add_line(self, line: IRCLine.ParsedLine):