Refactor module (un|re)loading

sopel/bot.py

@@ -90,6 +90,12 @@ def __init__(self, config, daemon=False):
         self._cap_reqs = dict()
         """A dictionary of capability names to a list of requests"""
 
+        self._modules = dict()
+        """A dictionary of modules currently registered by the bot. sys.modules
+        doesn't work because although del <module> removes it from the
+        namespace, it remains in sys.modules. Thus, we need another way to keep
+        track of what has been (un)loaded."""
+
         self.privileges = dict()
         """A dictionary of channels to their users and privilege levels
 
@@ -177,53 +183,22 @@ def setup(self):
 
             try:
                 module, _ = sopel.loader.load_module(name, path, type_)
+                self.register_module(module)
             except Exception as e:
                 error_count = error_count + 1
                 filename, lineno = tools.get_raising_file_and_line()
                 rel_path = os.path.relpath(filename, os.path.dirname(__file__))
                 raising_stmt = "%s:%d" % (rel_path, lineno)
                 stderr("Error loading %s: %s (%s)" % (name, e, raising_stmt))
             else:
-                try:
-                    if hasattr(module, 'setup'):
-                        module.setup(self)
-                    relevant_parts = sopel.loader.clean_module(
-                        module, self.config)
-                except Exception as e:
-                    error_count = error_count + 1
-                    filename, lineno = tools.get_raising_file_and_line()
-                    rel_path = os.path.relpath(
-                        filename, os.path.dirname(__file__)
-                    )
-                    raising_stmt = "%s:%d" % (rel_path, lineno)
-                    stderr("Error in %s setup procedure: %s (%s)"
-                           % (name, e, raising_stmt))
-                else:
-                    self.register(*relevant_parts)
-                    success_count += 1
+                success_count += 1
 
         if len(modules) > 1:  # coretasks is counted
             stderr('\n\nRegistered %d modules,' % (success_count - 1))
             stderr('%d modules failed to load\n\n' % error_count)
         else:
             stderr("Warning: Couldn't load any modules")
 
-    def unregister(self, obj):
-        if not callable(obj):
-            return
-        if hasattr(obj, 'rule'):  # commands and intents have it added
-            for rule in obj.rule:
-                callb_list = self._callables[obj.priority][rule]
-                if obj in callb_list:
-                    callb_list.remove(obj)
-        if hasattr(obj, 'interval'):
-            # TODO this should somehow find the right job to remove, rather than
-            # clearing the entire queue. Issue #831
-            self.scheduler.clear_jobs()
-        if (getattr(obj, '__name__', None) == 'shutdown' and
-                    obj in self.shutdown_methods):
-            self.shutdown_methods.remove(obj)
-
     def register(self, callables, jobs, shutdowns, urls):
         # Append module's shutdown function to the bot's list of functions to
         # call on shutdown
@@ -252,6 +227,67 @@ def register(self, callables, jobs, shutdowns, urls):
         for func in urls:
             self.memory['url_callbacks'][func.url_regex] = func
 
+    def register_module(self, module):
+        try:
+            if hasattr(module, 'setup'):
+                module.setup(self)
+        except Exception as e:
+            raise RuntimeError("Error in setup procedure: %s" % e)
+        else:
+            relevant_parts = sopel.loader.clean_module(module, self.config)
+            self.register(*relevant_parts)
+            self._modules[module.__name__] = module
+
+    def unregister(self, callables, jobs, shutdowns, urls):
+        re_dotstar = re.compile('.*')
+
+        for shutdown in shutdowns:
+            try:
+                self.shutdown_methods.remove(shutdown)
+            except ValueError:
+                pass
+
+        for callbl in callables:
+            if hasattr(callbl, 'rule'):
+                for rule in callbl.rule:
+                    try:
+                        self._callables[callbl.priority][rule].remove(callbl)
+                    except ValueError:
+                        pass
+            else:
+                try:
+                    self._callables[callbl.priority][re_dotstar].remove(callbl)
+                except ValueError:
+                    pass
+
+            if hasattr(callbl, 'commands'):
+                module_name = callbl.__module__.rsplit('.', 1)[-1]
+                # TODO doc and make decorator for this. Not sure if this is how
+                # it should work yet, so not making it public for 6.0.
+                category = getattr(callbl, 'category', module_name)
+
+                if callbl.commands[0] in self._command_groups[category]:
+                    self._command_groups[category].remove(callbl.commands[0])
+
+            for command in callbl._docs:
+                self.doc.pop(command, None)
+
+        for func in jobs:
+            for interval in func.interval:
+                self.scheduler.del_job_by_params(interval, func)
+
+        for func in urls:
+            self.memory['url_callbacks'].pop(func.url_regex, None)
+
+    def unregister_module(self, module):
+        relevant_parts = sopel.loader.clean_module(module, self.config, modify=False)
+        self.unregister(*relevant_parts)
+
+        if hasattr(module, "teardown"):
+            module.teardown(self)
+
+        self._modules.pop(module.__name__, None)
+
     def part(self, channel, msg=None):
         """Part a channel."""
         self.write(['PART', channel], msg)

sopel/loader.py

@@ -2,16 +2,27 @@
 from __future__ import unicode_literals, absolute_import, print_function, division
 
 import imp
+import importlib
 import os.path
 import re
 import sys
+from types import ModuleType
 
 from sopel.tools import compile_rule, itervalues, get_command_regexp, get_nickname_command_regexp
 
 if sys.version_info.major >= 3:
     basestring = (str, bytes)
 
 
+try:
+    _reload = reload
+except NameError:
+    try:
+        _reload = importlib.reload
+    except AttributeError:
+        _reload = imp.reload
+
+
 def get_module_description(path):
     good_file = (os.path.isfile(path) and
                  path.endswith('.py') and not path.startswith('_'))
@@ -192,28 +203,86 @@ def load_module(name, path, type_):
         module = imp.load_module(name, None, path, ('', '', type_))
     else:
         raise TypeError('Unsupported module type')
+
     return module, os.path.getmtime(path)
 
 
 def is_triggerable(obj):
     return any(hasattr(obj, attr) for attr in ('rule', 'intents', 'commands', 'nickname_commands'))
 
 
-def clean_module(module, config):
+def clean_module(module, config, modify=True):
     callables = []
     shutdowns = []
     jobs = []
     urls = []
+
     for obj in itervalues(vars(module)):
         if callable(obj):
             if getattr(obj, '__name__', None) == 'shutdown':
                 shutdowns.append(obj)
             elif is_triggerable(obj):
-                clean_callable(obj, config)
+                if modify:
+                    clean_callable(obj, config)
                 callables.append(obj)
             elif hasattr(obj, 'interval'):
-                clean_callable(obj, config)
+                if modify:
+                    clean_callable(obj, config)
                 jobs.append(obj)
             elif hasattr(obj, 'url_regex'):
                 urls.append(obj)
+
     return callables, jobs, shutdowns, urls
+
+
+# https://github.com/thodnev/reload_all
+def reload_all(top_module, max_depth=20, raise_immediately=False,
+               pre_reload=None, reload_if=None):
+    '''
+    A reload function, which recursively traverses through
+    all submodules of top_module and reloads them from most-
+    nested to least-nested. Only modules containing __file__
+    attribute could be reloaded.
+
+    Returns a dict of not reloaded(due to errors) modules:
+      key = module, value = exception
+    Optional attribute max_depth defines maximum recursion
+    limit to avoid infinite loops while tracing
+    '''
+    # modules to reload: K=module, V=depth
+    for_reload = dict()
+
+    def trace_reload(module, depth):  # recursive
+        depth += 1
+
+        if isinstance(module, ModuleType) and depth < max_depth:
+            # check condition if provided
+            if reload_if is not None and not reload_if(module, depth):
+                return
+
+            # if module is deeper and could be reloaded
+            if for_reload.get(module, 0) < depth and hasattr(module, '__file__'):
+                for_reload[module] = depth
+
+            # trace through all attributes recursively
+            for attr in module.__dict__.values():
+                trace_reload(attr, depth)
+
+    # start tracing
+    trace_reload(top_module, 0)
+    reload_list = sorted(for_reload, reverse=True, key=lambda k: for_reload[k])
+    not_reloaded = dict()
+
+    if pre_reload is not None:
+        for module in reload_list:
+            pre_reload(module)
+
+    for module in reload_list:
+        try:
+            _reload(module)
+        except Exception:  # catch and write all errors
+            if raise_immediately:
+                raise
+            not_reloaded[module] = sys.exc_info()[0]
+
+    return not_reloaded