From 295f7a751c9f4917ddd6017c61945e7588b7000d Mon Sep 17 00:00:00 2001
From: jesopo <github@lolnerd.net>
Date: Sat, 11 May 2019 15:11:09 +0100
Subject: [PATCH] Refactor CAPs in to objects

---
 modules/line_handler/ircv3.py | 30 ++++++++++++++++++++++++------
 src/utils/irc/__init__.py     | 11 +++++++++++
 2 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/modules/line_handler/ircv3.py b/modules/line_handler/ircv3.py
index 077b6ec3..ea68c67b 100644
--- a/modules/line_handler/ircv3.py
+++ b/modules/line_handler/ircv3.py
@@ -1,13 +1,31 @@
 from src import utils
 
-CAPABILITIES = {"multi-prefix", "chghost", "invite-notify", "account-tag",
-    "account-notify", "extended-join", "away-notify", "userhost-in-names",
-    "draft/message-tags-0.2", "message-tags", "server-time", "cap-notify",
-    "batch", "draft/labeled-response", "draft/rename", "echo-message",
-    "draft/setname"}
+CAPABILITIES = [
+    utils.irc.Capability("multi-prefix"),
+    utils.irc.Capability("chghost"),
+    utils.irc.Capability("invite-notify"),
+    utils.irc.Capability("account-tag"),
+    utils.irc.Capability("account-notify"),
+    utils.irc.Capability("extended-join"),
+    utils.irc.Capability("away-notify"),
+    utils.irc.Capability("userhost-in-names"),
+    utils.irc.Capability("message-tags", "draft/message-tags-0.2"),
+    utils.irc.Capability("server-time"),
+    utils.irc.Capability("cap-notify"),
+    utils.irc.Capability("batch"),
+    utils.irc.Capability("echo-message"),
+    utils.irc.Capability(None, "draft/labeled-response"),
+    utils.irc.Capability(None, "draft/rename"),
+    utils.irc.Capability(None, "draft/setname")
+]
 
 def _match_caps(capabilities):
-    return set(capabilities) & CAPABILITIES
+    matched = []
+    for capability in CAPABILITIES:
+        available = capability.available(capabilities)
+        if available:
+            matched.append(available)
+    return matched
 
 def cap(events, event):
     capabilities = utils.parse.keyvalue(event["args"][-1])
diff --git a/src/utils/irc/__init__.py b/src/utils/irc/__init__.py
index 3e285995..15850cea 100644
--- a/src/utils/irc/__init__.py
+++ b/src/utils/irc/__init__.py
@@ -275,3 +275,14 @@ class IRCSendBatch(IRCBatch):
         self._add_line(utils.irc.protocol.message(target, message, tags))
     def notice(self, target: str, message: str, tags: dict={}):
         self._add_line(utils.irc.protocol.notice(target, message, tags))
+
+class Capability(object):
+    def __init__(self, name, draft_name=None):
+        self._caps = set([name, draft_name])
+        self._name = name
+        self._draft_name = draft_name
+    def available(self, capabilities: typing.List[str]) -> str:
+        match = list(set(capabilities)&self._caps)
+        return match[0] if match else None
+    def enabled(self, capability: str) -> bool:
+        return capability in self._caps