Hold on to tags and type for BATCHes

2 files changed, 15 insertions, 4 deletions

diff --git a/modules/line_handler.py b/modules/line_handler.py
index bac0245b..8b1dbf90 100644
--- a/modules/line_handler.py
+++ b/modules/line_handler.py
@@ -8,6 +8,7 @@ CAPABILITIES = {"multi-prefix", "chghost", "invite-notify", "account-tag",
     "account-notify", "extended-join", "away-notify", "userhost-in-names",
     "draft/message-tags-0.2", "draft/message-tags-0.3", "server-time",
     "cap-notify", "batch", "draft/labeled-response", "draft/rename"}
+LABELED_BATCH = ["draft/labeled-response", "labeled-response"]
 
 class Direction(enum.Enum):
     SEND = 0
@@ -38,7 +39,7 @@ class Module(ModuleManager.BaseModule):
         line = utils.irc.parse_line(event["line"])
         if "batch" in line.tags and line.tags["batch"] in event[
                 "server"].batches:
-            server.batches[tag["batch"]].append(line)
+            server.batches[tag["batch"]].lines.append(line)
         else:
             self._handle(event["server"], line)
 
@@ -554,12 +555,15 @@ class Module(ModuleManager.BaseModule):
     def batch(self, event):
         identifier = event["args"][0]
         modifier, identifier = identifier[0], identifier[1:]
+        batch_type = event["args"][1]
+
         if modifier == "+":
-            event["server"].batches[identifier] = []
+            event["server"].batches[identifier] = utils.irc.IRCBatch(identifier,
+                batch_type, event["tags"])
         else:
-            lines = event["server"].batches[identifier]
+            batch = event["server"].batches[identifier]
             del event["server"].batches[identifier]
-            for line in lines:
+            for line in batch.lines:
                 self._handle(line)
 
     # IRCv3 CHGHOST, a user's username and/or hostname has changed
diff --git a/src/utils/irc.py b/src/utils/irc.py
index 48f6800e..78b26d65 100644
--- a/src/utils/irc.py
+++ b/src/utils/irc.py
@@ -286,3 +286,10 @@ def parse_ctcp(s: str) -> typing.Optional[CTCPMessage]:
         return CTCPMessage(ctcp_command, ctcp_message)
 
     return None
+
+class IRCBatch(object):
+    def __init__(self, identifier: str, batch_type: str, tags: dict):
+        self.id = identifier
+        self.type = batch_type
+        self.tags = tags
+        self.lines = [] # type: typing.List[IRCParsedLine]