Add on_connect hooks and move nick/user/pass/cap ls command out of the core

Author: linuxdaemon

cloudbot/clients/irc.py

@@ -1,8 +1,8 @@
-from _ssl import PROTOCOL_SSLv23
 import asyncio
+import logging
 import re
 import ssl
-import logging
+from _ssl import PROTOCOL_SSLv23
 from ssl import SSLContext
 
 from cloudbot.client import Client
@@ -122,12 +122,11 @@ def connect(self):
         self._transport, self._protocol = yield from self.loop.create_connection(
             lambda: _IrcProtocol(self), host=self.server, port=self.port, ssl=self.ssl_context, **optional_params)
 
-        # send the cap ls, password, nick, and user
-        self.send("CAP LS 302")
-        self.set_pass(self.config["connection"].get("password"))
-        self.set_nick(self.nick)
-        self.cmd("USER", self.config.get('user', 'cloudbot'), "3", "*",
-                 self.config.get('realname', 'CloudBot - https://git.io/CloudBot'))
+        tasks = [
+            self.bot.plugin_manager.launch(hook, Event(bot=self.bot, conn=self, hook=hook))
+            for hook in self.bot.plugin_manager.connect_hooks
+        ]
+        yield from asyncio.gather(*tasks)
 
     def quit(self, reason=None):
         if self._quit:

cloudbot/hook.py

@@ -401,3 +401,21 @@ def _on_cap_ack_hook(func):
         return func
 
     return _on_cap_ack_hook
+
+
+def on_connect(param=None, **kwargs):
+    def _on_connect_hook(func):
+        hook = _get_hook(func, "on_connect")
+        if hook is None:
+            hook = _Hook(func, "on_connect")
+            _add_hook(func, hook)
+        hook._add_hook(kwargs)
+        return func
+
+    if callable(param):
+        return _on_connect_hook(param)
+    else:
+        return lambda func: _on_connect_hook(func)
+
+
+connect = on_connect

cloudbot/plugin.py

@@ -6,6 +6,7 @@
 import os
 import re
 from collections import defaultdict
+from operator import attrgetter
 
 import sqlalchemy
 
@@ -19,7 +20,7 @@ def find_hooks(parent, module):
     """
     :type parent: Plugin
     :type module: object
-    :rtype: (list[CommandHook], list[RegexHook], list[RawHook], list[SieveHook], List[EventHook], List[PeriodicHook], list[OnStartHook], List[OnStopHook], list[OnCapAckHook], list[OnCapAvailableHook])
+    :rtype: (list[CommandHook], list[RegexHook], list[RawHook], list[SieveHook], List[EventHook], List[PeriodicHook], list[OnStartHook], List[OnStopHook], list[OnCapAckHook], list[OnCapAvailableHook], list[OnConnectHook])
     """
     # set the loaded flag
     module._cloudbot_loaded = True
@@ -33,9 +34,10 @@ def find_hooks(parent, module):
     on_stop = []
     on_cap_ack = []
     on_cap_available = []
+    on_connect = []
     type_lists = {"command": command, "regex": regex, "irc_raw": raw, "sieve": sieve, "event": event,
                   "periodic": periodic, "on_start": on_start, "on_stop": on_stop, "on_cap_ack": on_cap_ack,
-                  "on_cap_available": on_cap_available}
+                  "on_cap_available": on_cap_available, "on_connect": on_connect}
     for name, func in module.__dict__.items():
         if hasattr(func, "_cloudbot_hook"):
             # if it has cloudbot hook
@@ -47,7 +49,7 @@ def find_hooks(parent, module):
             # delete the hook to free memory
             del func._cloudbot_hook
 
-    return command, regex, raw, sieve, event, periodic, on_start, on_stop, on_cap_ack, on_cap_available
+    return command, regex, raw, sieve, event, periodic, on_start, on_stop, on_cap_ack, on_cap_available, on_connect
 
 
 def find_tables(code):
@@ -103,6 +105,7 @@ def __init__(self, bot):
         self.regex_hooks = []
         self.sieves = []
         self.cap_hooks = {"on_available": defaultdict(list), "on_ack": defaultdict(list)}
+        self.connect_hooks = []
         self._hook_waiting_queues = {}
 
     @asyncio.coroutine
@@ -242,8 +245,14 @@ def load_plugin(self, path):
             self.sieves.append(sieve_hook)
             self._log_hook(sieve_hook)
 
+        # register connect hooks
+        for connect_hook in plugin.connect_hooks:
+            self.connect_hooks.append(connect_hook)
+            self._log_hook(connect_hook)
+
         # sort sieve hooks by priority
         self.sieves.sort(key=lambda x: x.priority)
+        self.connect_hooks.sort(key=attrgetter("priority"))
 
         # we don't need this anymore
         del plugin.run_on_start
@@ -325,6 +334,10 @@ def unload_plugin(self, path):
         for sieve_hook in plugin.sieves:
             self.sieves.remove(sieve_hook)
 
+        # unregister connect hooks
+        for connect_hook in plugin.connect_hooks:
+            self.connect_hooks.remove(connect_hook)
+
         # Run on_stop hooks
         for on_stop_hook in plugin.run_on_stop:
             event = Event(bot=self.bot, hook=on_stop_hook)
@@ -552,8 +565,11 @@ def __init__(self, filepath, filename, title, code):
         self.file_path = filepath
         self.file_name = filename
         self.title = title
-        self.commands, self.regexes, self.raw_hooks, self.sieves, self.events, self.periodic, self.run_on_start, self.run_on_stop, self.on_cap_ack, self.on_cap_available = find_hooks(
-            self, code)
+        hooks = find_hooks(self, code)
+        self.commands, self.regexes, self.raw_hooks, *hooks = hooks
+        self.sieves, self.events, self.periodic, *hooks = hooks
+        self.run_on_start, self.run_on_stop, self.on_cap_ack, *hooks = hooks
+        self.on_cap_available, self.connect_hooks, *hooks = hooks
         # we need to find tables for each plugin so that they can be unloaded from the global metadata when the
         # plugin is reloaded
         self.tables = find_tables(code)
@@ -830,6 +846,23 @@ def __init__(self, plugin, base_hook):
         super().__init__("ack", plugin, base_hook)
 
 
+class OnConnectHook(Hook):
+    def __init__(self, plugin, sieve_hook):
+        """
+        :type plugin: Plugin
+        :type sieve_hook: cloudbot.util.hook._Hook
+        """
+
+        self.priority = sieve_hook.kwargs.pop("priority", 100)
+        super().__init__("on_connect", plugin, sieve_hook)
+
+    def __repr__(self):
+        return "{name}[{base!r}]".format(name=self.type, base=super())
+
+    def __str__(self):
+        return "{name} {func} from {file}".format(name=self.type, func=self.function_name, file=self.plugin.file_name)
+
+
 _hook_name_to_plugin = {
     "command": CommandHook,
     "regex": RegexHook,
@@ -841,4 +874,5 @@ def __init__(self, plugin, base_hook):
     "on_stop": OnStopHook,
     "on_cap_available": OnCapAvaliableHook,
     "on_cap_ack": OnCapAckHook,
+    "on_connect": OnConnectHook,
 }

plugins/cap.py

@@ -8,6 +8,12 @@
 logger = logging.getLogger("cloudbot")
 
 
+@asyncio.coroutine
+@hook.connect(priority=-10)
+def send_cap_ls(conn):
+    conn.cmd("CAP", "LS", "302")
+
+
 @asyncio.coroutine
 def handle_available_caps(conn, caplist, event, irc_paramlist, bot):
     available_caps = conn.memory.setdefault("available_caps", set())

plugins/core_connect.py

@@ -0,0 +1,19 @@
+from cloudbot import hook
+
+
+@hook.connect(priority=0)
+def conn_pass(conn):
+    conn.set_pass(conn.config["connection"].get("password"))
+
+
+@hook.connect(priority=10)
+def conn_nick(conn):
+    conn.set_nick(conn.nick)
+
+
+@hook.connect(priority=20)
+def conn_user(conn):
+    conn.cmd(
+        "USER", conn.config.get('user', 'cloudbot'), "3", "*",
+        conn.config.get('realname', 'CloudBot - https://git.io/CloudBot')
+    )