[commands] Add Bot.reload_extension for atomic loading.

Also do atomic loading in Bot.load_extension

discord/ext/commands/bot.py

@@ -523,6 +523,65 @@ class BotBase(GroupMixin):
 
     # extensions
 
+    def _remove_module_references(self, name):
+        # find all references to the module
+        # remove the cogs registered from the module
+        for cogname, cog in self._cogs.copy().items():
+            if _is_submodule(name, cog.__module__):
+                self.remove_cog(cogname)
+
+        # remove all the commands from the module
+        for cmd in self.all_commands.copy().values():
+            if cmd.module is not None and _is_submodule(name, cmd.module):
+                if isinstance(cmd, GroupMixin):
+                    cmd.recursively_remove_all_commands()
+                self.remove_command(cmd.name)
+
+        # remove all the listeners from the module
+        for event_list in self.extra_events.copy().values():
+            remove = []
+            for index, event in enumerate(event_list):
+                if event.__module__ is not None and _is_submodule(name, event.__module__):
+                    remove.append(index)
+
+            for index in reversed(remove):
+                del event_list[index]
+
+    def _call_module_finalizers(self, lib, key):
+        try:
+            func = getattr(lib, 'teardown')
+        except AttributeError:
+            pass
+        else:
+            try:
+                func(self)
+            except Exception:
+                pass
+        finally:
+            self._extensions.pop(key, None)
+            sys.modules.pop(key, None)
+            name = lib.__name__
+            for module in list(sys.modules.keys()):
+                if _is_submodule(name, module):
+                    del sys.modules[module]
+
+    def _load_from_module_spec(self, lib, key):
+        # precondition: key not in self._extensions
+        try:
+            setup = getattr(lib, 'setup')
+        except AttributeError:
+            del sys.modules[key]
+            raise discord.ClientException('extension {!r} ({!r}) does not have a setup function.'.format(key, lib))
+
+        try:
+            setup(self)
+        except Exception:
+            self._remove_module_references(lib.__name__)
+            self._call_module_finalizers(lib, key)
+            raise
+        else:
+            self._extensions[key] = lib
+
     def load_extension(self, name):
         """Loads an extension.
 
@@ -546,19 +605,16 @@ class BotBase(GroupMixin):
             The extension does not have a setup function.
         ImportError
             The extension could not be imported.
+        Exception
+            Any other exception raised by the extension will be raised back
+            to the caller.
         """
 
         if name in self._extensions:
             return
 
         lib = importlib.import_module(name)
-        if not hasattr(lib, 'setup'):
+        self._load_from_module_spec(lib, name)
-            del lib
-            del sys.modules[name]
-            raise discord.ClientException('extension does not have a setup function')
-
-        lib.setup(self)
-        self._extensions[name] = lib
 
     def unload_extension(self, name):
         """Unloads an extension.
@@ -583,49 +639,56 @@ class BotBase(GroupMixin):
         if lib is None:
             return
 
-        lib_name = lib.__name__
+        self._remove_module_references(lib.__name__)
+        self._call_module_finalizers(lib, name)
 
-        # find all references to the module
+    def reload_extension(self, name):
+        """Atomically reloads an extension.
 
-        # remove the cogs registered from the module
+        This replaces the extension with the same extension, only refreshed. This is
-        for cogname, cog in self._cogs.copy().items():
+        equivalent to a :meth:`unload_extension` followed by a :meth:`load_extension`
-            if _is_submodule(lib_name, cog.__module__):
+        except done in an atomic way. That is, if an operation fails mid-reload then
-                self.remove_cog(cogname)
+        the bot will roll-back to the prior working state.
 
-        # remove all the commands from the module
+        Parameters
-        for cmd in self.all_commands.copy().values():
+        ------------
-            if cmd.module is not None and _is_submodule(lib_name, cmd.module):
+        name: :class:`str`
-                if isinstance(cmd, GroupMixin):
+            The extension name to reload. It must be dot separated like
-                    cmd.recursively_remove_all_commands()
+            regular Python imports if accessing a sub-module. e.g.
-                self.remove_command(cmd.name)
+            ``foo.test`` if you want to import ``foo/test.py``.
 
-        # remove all the listeners from the module
+        Raises
-        for event_list in self.extra_events.copy().values():
+        -------
-            remove = []
+        Exception
-            for index, event in enumerate(event_list):
+            Any exception raised by the extension will be raised back
-                if event.__module__ is not None and _is_submodule(lib_name, event.__module__):
+            to the caller.
-                    remove.append(index)
+        """
 
-            for index in reversed(remove):
+        lib = self._extensions.get(name)
-                del event_list[index]
+        if lib is None:
+            return
+
+        # get the previous module states from sys modules
+        modules = {
+            name: module
+            for name, module in sys.modules.items()
+            if _is_submodule(lib.__name__, name)
+        }
 
         try:
-            func = getattr(lib, 'teardown')
+            # Unload and then load the module...
-        except AttributeError:
+            self._remove_module_references(lib.__name__)
-            pass
+            self._call_module_finalizers(lib, name)
-        else:
+            self.load_extension(name)
-            try:
+        except Exception as e:
-                func(self)
+            # if the load failed, the remnants should have been
-            except Exception:
+            # cleaned from the load_extension function call
-                pass
+            # so let's load it from our old compiled library.
-        finally:
+            self._load_from_module_spec(lib, name)
-            # finally remove the import..
+
-            del lib
+            # revert sys.modules back to normal and raise back to caller
-            del self._extensions[name]
+            sys.modules.update(modules)
-            del sys.modules[name]
+            raise
-            for module in list(sys.modules.keys()):
-                if _is_submodule(lib_name, module):
-                    del sys.modules[module]
 
     @property
     def extensions(self):

docs/ext/commands/extensions.rst

@@ -41,14 +41,13 @@ In this example we define a simple command, and when the extension is loaded thi
 Reloading
 -----------
 
-The act of reloading an extension is actually quite simple -- it is as simple as unloading it and then reloading it.
+When you make a change to the extension and want to reload the references, the library comes with a function to do this for you, :meth:`Bot.reload_extension`.
 
 .. code-block:: python3
 
-    >>> bot.unload_extension('hello')
+    >>> bot.reload_extension('hello')
-    >>> bot.load_extension('hello')
 
-Once we remove and load the extension, any changes that we did will be applied upon load. This is useful if we want to add or remove functionality without restarting our bot.
+Once the extension reloads, any changes that we did will be applied. This is useful if we want to add or remove functionality without restarting our bot. If an error occurred during the reloading process, the bot will pretend as if the reload never happened.
 
 Cleaning Up
 -------------