Rework extension loading (#1423)

Co-authored-by: Lala Sabathil <[email protected]>

Author: baronkobama

discord/cog.py

@@ -26,6 +26,8 @@
 
 import importlib
 import inspect
+import os
+import pathlib
 import sys
 import types
 from typing import (
@@ -40,6 +42,7 @@
     Tuple,
     Type,
     TypeVar,
+    Union,
 )
 
 import discord.utils
@@ -739,7 +742,14 @@ def _resolve_name(self, name: str, package: Optional[str]) -> str:
         except ImportError:
             raise errors.ExtensionNotFound(name)
 
-    def load_extension(self, name: str, *, package: Optional[str] = None) -> None:
+    def load_extension(
+        self,
+        name: str,
+        *,
+        package: Optional[str] = None,
+        recursive: bool = False,
+        store: bool = True,
+    ) -> Optional[Union[Dict[str, Union[Exception, bool]], List[str]]]:
         """Loads an extension.
 
         An extension is a python module that contains commands, cogs, or
@@ -749,21 +759,41 @@ def load_extension(self, name: str, *, package: Optional[str] = None) -> None:
         the entry point on what to do when the extension is loaded. This entry
         point must have a single argument, the ``bot``.
 
+        The extension passed can either be the direct name of a file within
+        the current working directory or a folder that contains multiple extensions.
+
         Parameters
-        ------------
+        -----------
         name: :class:`str`
-            The extension name to load. It must be dot separated like
-            regular Python imports if accessing a sub-module. e.g.
+            The extension or folder name to load. It must be dot separated
+            like regular Python imports if accessing a sub-module. e.g.
             ``foo.test`` if you want to import ``foo/test.py``.
         package: Optional[:class:`str`]
             The package name to resolve relative imports with.
-            This is required when loading an extension using a relative path, e.g ``.foo.test``.
+            This is required when loading an extension using a relative
+            path, e.g ``.foo.test``.
             Defaults to ``None``.
 
             .. versionadded:: 1.7
+        recursive: Optional[:class:`bool`]
+            If subdirectories under the given head directory should be
+            recursively loaded.
+            Defaults to ``False``.
+
+            .. versionadded:: 2.0
+        store: Optional[:class:`bool`]
+            If exceptions should be stored or raised. If set to ``True``,
+            all exceptions encountered will be stored in a returned dictionary
+            as a load status. If set to ``False``, if any exceptions are
+            encountered they will be raised and the bot will be closed.
+            If no exceptions are encountered, a list of loaded
+            extension names will be returned.
+            Defaults to ``True``.
+
+            .. versionadded:: 2.0
 
         Raises
-        --------
+        -------
         ExtensionNotFound
             The extension could not be imported.
             This is also raised if the name of the extension could not
@@ -774,17 +804,133 @@ def load_extension(self, name: str, *, package: Optional[str] = None) -> None:
             The extension does not have a setup function.
         ExtensionFailed
             The extension or its setup function had an execution error.
+
+        Returns
+        --------
+        Optional[Union[Dict[:class:`str`, Union[:exc:`errors.ExtensionError`, :class:`bool`]], List[:class:`str`]]]
+            If the store parameter is set to ``True``, a dictionary will be returned that
+            contains keys to represent the loaded extension names. The values bound to
+            each key can either be an exception that occurred when loading that extension
+            or a ``True`` boolean representing a successful load. If the store parameter
+            is set to ``False``, either a list containing a list of loaded extensions or
+            nothing due to an encountered exception.
         """
 
         name = self._resolve_name(name, package)
+
         if name in self.__extensions:
-            raise errors.ExtensionAlreadyLoaded(name)
+            exc = errors.ExtensionAlreadyLoaded(name)
+            final_out = {name: exc} if store else exc
+        # This indicates that there is neither an extension nor folder here
+        elif (spec := importlib.util.find_spec(name)) is None:
+            exc = errors.ExtensionNotFound(name)
+            final_out = {name: exc} if store else exc
+        # This indicates we've found an extension file to load, and we need to store any exceptions
+        elif spec.has_location and store:
+            try:
+                self._load_from_module_spec(spec, name)
+            except Exception as exc:
+                final_out = {name: exc}
+            else:
+                final_out = {name: True}
+        # This indicates we've found an extension file to load, and any encountered exceptions can be raised
+        elif spec.has_location:
+            self._load_from_module_spec(spec, name)
+            final_out = [name]
+        # This indicates we've been given a folder because the ModuleSpec exists but is not a file
+        else:
+            # Split the directory path and join it to get an os-native Path object
+            path = pathlib.Path(os.path.join(*name.split(".")))
+            glob = path.rglob if recursive else path.glob
+            final_out = {} if store else []
+
+            # Glob all files with a pattern to gather all .py files that don't start with _
+            for ext_file in glob("[!_]*.py"):
+                # Gets all parts leading to the directory minus the file name
+                parts = list(ext_file.parts[:-1])
+                # Gets the file name without the extension
+                parts.append(ext_file.stem)
+                loaded = self.load_extension(".".join(parts))
+                final_out.update(loaded) if store else final_out.extend(loaded)
+
+        if isinstance(final_out, Exception):
+            raise final_out
+        else:
+            return final_out
 
-        spec = importlib.util.find_spec(name)
-        if spec is None:
-            raise errors.ExtensionNotFound(name)
+    def load_extensions(
+        self,
+        *names: str,
+        package: Optional[str] = None,
+        recursive: bool = False,
+        store: bool = True,
+    ) -> Optional[Union[Dict[str, Union[Exception, bool]], List[str]]]:
+        """Loads multiple extensions at once.
+
+        This method simplifies the process of loading multiple
+        extensions by handling the looping of ``load_extension``.
+
+        Parameters
+        -----------
+        names: :class:`str`
+           The extension or folder names to load. It must be dot separated
+           like regular Python imports if accessing a sub-module. e.g.
+           ``foo.test`` if you want to import ``foo/test.py``.
+        package: Optional[:class:`str`]
+            The package name to resolve relative imports with.
+            This is required when loading an extension using a relative
+            path, e.g ``.foo.test``.
+            Defaults to ``None``.
+
+            .. versionadded:: 1.7
+        recursive: Optional[:class:`bool`]
+            If subdirectories under the given head directory should be
+            recursively loaded.
+            Defaults to ``False``.
+
+            .. versionadded:: 2.0
+        store: Optional[:class:`bool`]
+            If exceptions should be stored or raised. If set to ``True``,
+            all exceptions encountered will be stored in a returned dictionary
+            as a load status. If set to ``False``, if any exceptions are
+            encountered they will be raised and the bot will be closed.
+            If no exceptions are encountered, a list of loaded
+            extension names will be returned.
+            Defaults to ``True``.
+
+            .. versionadded:: 2.0
+
+        Raises
+        --------
+        ExtensionNotFound
+            A given extension could not be imported.
+            This is also raised if the name of the extension could not
+            be resolved using the provided ``package`` parameter.
+        ExtensionAlreadyLoaded
+            A given extension is already loaded.
+        NoEntryPointError
+            A given extension does not have a setup function.
+        ExtensionFailed
+            A given extension or its setup function had an execution error.
+
+        Returns
+        --------
+        Optional[Union[Dict[:class:`str`, Union[:exc:`errors.ExtensionError`, :class:`bool`]], List[:class:`str`]]]
+            If the store parameter is set to ``True``, a dictionary will be returned that
+            contains keys to represent the loaded extension names. The values bound to
+            each key can either be an exception that occurred when loading that extension
+            or a ``True`` boolean representing a successful load. If the store parameter
+            is set to ``False``, either a list containing names of loaded extensions or
+            nothing due to an encountered exception.
+        """
+
+        loaded_extensions = {} if store else []
+
+        for ext_path in names:
+            loaded = self.load_extension(ext_path, package=package, recursive=recursive, store=store)
+            loaded_extensions.update(loaded) if store else loaded_extensions.extend(loaded)
 
-        self._load_from_module_spec(spec, name)
+        return loaded_extensions
 
     def unload_extension(self, name: str, *, package: Optional[str] = None) -> None:
         """Unloads an extension.