Add utils.irc.MessageTag and use it for both draft and ratified msgid

2 files changed, 10 insertions, 1 deletions

diff --git a/modules/commands/__init__.py b/modules/commands/__init__.py
index ee90ec24..ff0be268 100644
--- a/modules/commands/__init__.py
+++ b/modules/commands/__init__.py
@@ -7,6 +7,8 @@ COMMAND_METHODS = ["PRIVMSG", "NOTICE"]
 
 REGEX_ARG_NUMBER = re.compile(r"\$(\d+)(-?)")
 
+MSGID_TAG = utils.irc.MessageTag("msgid", "draft/msgid")
+
 def _command_method_validate(s):
     if s.upper() in COMMAND_METHODS:
         return s.upper()
@@ -133,7 +135,7 @@ class Module(ModuleManager.BaseModule):
             if not module_name and hasattr(hook.function, "__self__"):
                 module_name = hook.function.__self__._name
 
-            msgid = event["tags"].get("draft/msgid", None)
+            msgid = MSGID_TAG.get_value(event["tags"])
             statusmsg = "".join(event.get("statusmsg", []))
             stdout = outs.StdOut(event["server"], module_name, target, msgid,
                 statusmsg)
diff --git a/src/utils/irc/__init__.py b/src/utils/irc/__init__.py
index 2e1ef6e4..8b4debfe 100644
--- a/src/utils/irc/__init__.py
+++ b/src/utils/irc/__init__.py
@@ -295,3 +295,10 @@ class Capability(object):
             callback()
     def nak(self):
         pass
+
+def MessageTag(object):
+    def __init__(self, name, draft_name=None):
+        self._names = set([name, draft_name])
+    def get_value(self, tags: typing.Dict[str, str]) -> typing.Optional[str]:
+        key = list(set(tags.keys())&self._names)
+        return tags[key[0]] if key else None