Support hooking functions in modules with @Utils.hook

3 files changed, 29 insertions, 4 deletions

diff --git a/ModuleManager.py b/ModuleManager.py
index a67cacad..29248a36 100644
--- a/ModuleManager.py
+++ b/ModuleManager.py
@@ -1,5 +1,7 @@
 import glob, imp, inspect, os, sys, uuid
 
+BITBOT_HOOKS_MAGIC = "__bitbot_hooks"
+
 class ModuleManager(object):
     def __init__(self, bot, events, exports, directory="modules"):
         self.bot = bot
@@ -51,10 +53,22 @@ class ModuleManager(object):
             raise ImportError("module '%s' has a Module attribute but it is not a class.")
 
         context = str(uuid.uuid4())
-        module_object = module.Module(self.bot, self.events.new_context(
-            context), self.exports.new_context(context))
+        context_events = self.events.new_context(context)
+        context_exports = self.exports.new_context(context)
+        module_object = module.Module(self.bot, context_events,
+            context_exports)
+
         if not hasattr(module_object, "_name"):
             module_object._name = name.title()
+        for attribute_name in dir(module_object):
+            attribute = getattr(module_object, attribute_name)
+            if inspect.ismethod(attribute) and hasattr(attribute,
+                    BITBOT_HOOKS_MAGIC):
+                hooks = getattr(attribute, BITBOT_HOOKS_MAGIC)
+                for hook in hooks:
+                    context_events.on(hook["event"]).hook(attribute,
+                        **hook["kwargs"])
+
         module_object._context = context
         module_object._import_name = name
 
diff --git a/Utils.py b/Utils.py
index b9a442dc..e2da0c43 100644
--- a/Utils.py
+++ b/Utils.py
@@ -1,6 +1,7 @@
 import json, re, traceback, urllib.request, urllib.parse, urllib.error, ssl
 import string
 import bs4
+import ModuleManager
 
 USER_AGENT = ("Mozilla/5.0 (Windows NT 6.1; WOW64) AppleWebKit/537.36 "
     "(KHTML, like Gecko) Chrome/49.0.2623.87 Safari/537.36")
@@ -275,3 +276,12 @@ def get_closest_setting(event, setting, default=None):
 
 def prevent_highlight(nickname):
     return nickname[0]+"\u200d"+nickname[1:]
+
+def hook(event, **kwargs):
+    def _hook_func(func):
+        if not hasattr(func, ModuleManager.BITBOT_HOOKS_MAGIC):
+            setattr(func, ModuleManager.BITBOT_HOOKS_MAGIC, [])
+        getattr(func, ModuleManager.BITBOT_HOOKS_MAGIC).append(
+            {"event": event, "kwargs": kwargs})
+        return func
+    return _hook_func
diff --git a/modules/pong.py b/modules/pong.py
index 72987542..d67da431 100644
--- a/modules/pong.py
+++ b/modules/pong.py
@@ -1,8 +1,9 @@
-
+import Utils
 
 class Module(object):
     def __init__(self, bot, events, exports):
-        events.on("received.command.ping").hook(self.pong, help="Ping pong!")
+        pass
 
+    @Utils.hook("received.command.ping", help="Ping pong!")
     def pong(self, event):
         event["stdout"].write("Pong!")