refactor: Bot.load_extensions rework

changelog/796.deprecate.rst

@@ -0,0 +1 @@
+:func:`disnake.utils.search_directory` will be removed in a future version, in favor of :func:`disnake.ext.commands.Bot.find_extensions` which is the most common usecase and is more consistent.

changelog/796.feature.rst

@@ -0,0 +1,3 @@
+|commands| Improve :func:`Bot.load_extensions <ext.commands.Bot.load_extensions>`, add :func:`Bot.find_extensions <ext.commands.Bot.find_extensions>`.
+- Better support for more complex extension hierarchies.
+- New ``package``, ``ignore``, and ``load_callback`` parameters.

disnake/ext/commands/common_bot_base.py

@@ -16,6 +16,7 @@
     Callable,
     Dict,
     Generic,
+    Iterable,
     List,
     Mapping,
     Optional,
@@ -423,7 +424,7 @@ def _load_from_module_spec(self, spec: importlib.machinery.ModuleSpec, key: str)
     def _resolve_name(self, name: str, package: Optional[str]) -> str:
         try:
             return importlib.util.resolve_name(name, package)
-        except ImportError:
+        except (ValueError, ImportError):  # 3.8 raises ValueError instead of ImportError
             raise errors.ExtensionNotFound(name)
 
     def load_extension(self, name: str, *, package: Optional[str] = None) -> None:
@@ -574,18 +575,151 @@ def reload_extension(self, name: str, *, package: Optional[str] = None) -> None:
             sys.modules.update(modules)
             raise
 
-    def load_extensions(self, path: str) -> None:
-        """Loads all extensions in a directory.
+    def find_extensions(
+        self,
+        root_module: str,
+        *,
+        package: Optional[str] = None,
+        ignore: Optional[Union[Iterable[str], Callable[[str], bool]]] = None,
+    ) -> List[str]:
+        """
+        Finds all extensions in a given module, also traversing into sub-packages.
+
+        See :ref:`ext_commands_extensions_load` for details on how packages are found.
+
+        .. versionadded:: 2.7
+
+        .. note::
+            This imports all *packages* (not all modules) in the given path(s)
+            to access the ``__path__`` attribute for finding submodules,
+            unless they are filtered by the ``ignore`` parameter.
+
+        Parameters
+        ----------
+        root_module: :class:`str`
+            The module/package name to search in, for example ``cogs.admin``.
+            Also supports paths in the current working directory.
+        package: Optional[:class:`str`]
+            The package name to resolve relative imports with.
+            This is required when ``root_module`` is a relative module name, e.g ``.cogs.admin``.
+            Defaults to ``None``.
+        ignore: Optional[Union[Iterable[:class:`str`], Callable[[:class:`str`], :class:`bool`]]]
+            An iterable of module names to ignore, or a callable that's used for ignoring
+            modules (where the callable returning ``True`` results in the module being ignored).
+            Defaults to ``None``, i.e. no modules are ignored.
+
+            If it's an iterable, module names that start with any of the given strings will be ignored.
+
+        Raises
+        ------
+        ExtensionError
+            The given root module could not be found,
+            or the name of the root module could not be resolved using the provided ``package`` parameter.
+        ValueError
+            ``root_module`` is a path and outside of the cwd.
+        TypeError
+            The ``ignore`` parameter is of an invalid type.
+        ImportError
+            A package couldn't be imported.
+
+        Returns
+        -------
+        List[:class:`str`]
+            The list of full extension names.
+        """
+        if "/" in root_module or "\\" in root_module:
+            path = os.path.relpath(root_module)
+            if ".." in path:
+                raise ValueError(
+                    "Paths outside the cwd are not supported. Try using the module name instead."
+                )
+            root_module = path.replace(os.sep, ".")
+
+        # `find_spec` already calls `resolve_name`, but we want our custom error handling here
+        root_module = self._resolve_name(root_module, package)
+
+        if not (spec := importlib.util.find_spec(root_module)):
+            raise errors.ExtensionError(
+                f"Unable to find root module '{root_module}'", name=root_module
+            )
+
+        if not (paths := spec.submodule_search_locations):
+            raise errors.ExtensionError(
+                f"Module '{root_module}' is not a package", name=root_module
+            )
+
+        return list(disnake.utils._walk_modules(paths, prefix=f"{spec.name}.", ignore=ignore))
+
+    def load_extensions(
+        self,
+        root_module: str,
+        *,
+        package: Optional[str] = None,
+        ignore: Optional[Union[Iterable[str], Callable[[str], bool]]] = None,
+        load_callback: Optional[Callable[[str], None]] = None,
+    ) -> Union[List[str], List[Union[str, errors.ExtensionError]]]:
+        """
+        Loads all extensions in a given module, also traversing into sub-packages.
+
+        See :func:`find_extensions` for details.
 
         .. versionadded:: 2.4
 
+        .. versionchanged:: 2.7
+            Now accepts a module name instead of a filesystem path.
+            Improved package traversal, adding support for more complex extensions
+            with ``__init__.py`` files.
+            Also added ``package``, ``ignore``, and ``load_callback`` parameters.
+
+        .. note::
+            For further customization, you may use :func:`find_extensions`:
+
+            .. code-block:: python3
+
+                for extension_name in bot.find_extensions(...):
+                    ... # custom logic
+                    bot.load_extension(extension_name)
+
         Parameters
         ----------
-        path: :class:`str`
-            The path to search for extensions
+        root_module: :class:`str`
+            See :func:`find_extensions`.
+        package: Optional[:class:`str`]
+            See :func:`find_extensions`.
+        ignore: Optional[Union[Iterable[:class:`str`], Callable[[:class:`str`], :class:`bool`]]]
+            See :func:`find_extensions`.
+        load_callback: Optional[Callable[[:class:`str`], None]]
+            A callback that gets invoked with the extension name when each extension gets loaded.
+
+        Raises
+        ------
+        ExtensionError
+            The given root module could not be found,
+            or the name of the root module could not be resolved using the provided ``package`` parameter.
+            Other extension-related errors may also be raised
+            as this method calls :func:`load_extension` on all found extensions.
+            See :func:`load_extension` for further details on raised exceptions.
+        ValueError
+            ``root_module`` is a path and outside of the cwd.
+        TypeError
+            The ``ignore`` parameter is of an invalid type.
+        ImportError
+            A package (not module) couldn't be imported.
+
+        Returns
+        -------
+        List[:class:`str`]
+            The list of module names that have been loaded.
         """
-        for extension in disnake.utils.search_directory(path):
-            self.load_extension(extension)
+        ret: List[str] = []
+
+        for ext_name in self.find_extensions(root_module, package=package, ignore=ignore):
+            self.load_extension(ext_name)
+            ret.append(ext_name)
+            if load_callback:
+                load_callback(ext_name)
+
+        return ret
 
     @property
     def extensions(self) -> Mapping[str, types.ModuleType]:

disnake/utils.py

@@ -6,6 +6,7 @@
 import asyncio
 import datetime
 import functools
+import importlib
 import json
 import os
 import pkgutil
@@ -1282,9 +1283,12 @@ def format_dt(dt: Union[datetime.datetime, float], /, style: TimestampStyle = "f
     return f"<t:{int(dt)}:{style}>"
 
 
+@deprecated("disnake.ext.commands.Bot.find_extensions")
 def search_directory(path: str) -> Iterator[str]:
     """Walk through a directory and yield all modules.
 
+    .. deprecated:: 2.7
+
     Parameters
     ----------
     path: :class:`str`
@@ -1318,6 +1322,47 @@ def search_directory(path: str) -> Iterator[str]:
             yield prefix + name
 
 
+# this is similar to pkgutil.walk_packages, but with a few modifications
+def _walk_modules(
+    paths: Iterable[str],
+    prefix: str = "",
+    ignore: Optional[Union[Iterable[str], Callable[[str], bool]]] = None,
+) -> Iterator[str]:
+    if isinstance(ignore, str):
+        raise TypeError("`ignore` must be an iterable of strings or a callable")
+
+    if isinstance(ignore, Iterable):
+        ignore_tup = tuple(ignore)
+        ignore = lambda path: path.startswith(ignore_tup)
+    # else, it's already a callable or None
+
+    seen: Set[str] = set()
+
+    for _, name, ispkg in pkgutil.iter_modules(paths, prefix):
+        if ignore and ignore(name):
+            continue
+
+        if ispkg:
+            mod = importlib.import_module(name)
+
+            # if this module is a package but also has a `setup` function,
+            # yield it and don't look for other files in this module
+            if hasattr(mod, "setup"):
+                yield name
+                continue
+
+            sub_paths: List[str] = []
+            for p in mod.__path__ or []:
+                if p not in seen:
+                    seen.add(p)
+                    sub_paths.append(p)
+
+            if sub_paths:
+                yield from _walk_modules(sub_paths, prefix=f"{name}.", ignore=ignore)
+        else:
+            yield name
+
+
 def as_valid_locale(locale: str) -> Optional[str]:
     """
     Converts the provided locale name to a name that is valid for use with the API,

docs/ext/commands/extensions.rst

@@ -64,3 +64,60 @@ Although rare, sometimes an extension needs to clean-up or know when it's being
 
     def teardown(bot):
         print('I am being unloaded!')
+
+.. _ext_commands_extensions_load:
+
+Loading multiple extensions
+-----------------------------
+
+Commonly, you might have a package/folder that contains several modules.
+Instead of manually loading them one by one, you can use :meth:`.Bot.load_extensions` to load the entire package in one sweep.
+
+Consider the following directory structure:
+
+.. code-block::
+
+    my_bot/
+    ├─── cogs/
+    │    ├─── admin.py
+    │    ├─── fun.py
+    │    └─── other_complex_thing/
+    │         ├─── __init__.py (contains setup)
+    │         ├─── data.py
+    │         └─── models.py
+    └─── main.py
+
+Now, you could call :meth:`.Bot.load_extension` separately on ``cogs.admin``, ``cogs.fun``, and ``cogs.other_complex_thing``;
+however, if you add a new extension, you'd need to once again load it separately.
+
+Instead, you can use ``bot.load_extensions("my_bot.cogs")`` (or ``.load_extensions(".cogs", package=__package__)``)
+to load all of them automatically, without any extra work required.
+
+Customization
++++++++++++++++
+
+To adjust the loading process, for example to handle exceptions that may occur, use :meth:`.Bot.find_extensions`.
+This is also what :meth:`.Bot.load_extensions` uses internally.
+
+As an example, one could load extensions like this:
+
+.. code-block:: python3
+
+    for extension in bot.find_extensions("my_bot.cogs"):
+        try:
+            bot.load_extension(extension)
+        except commands.ExtensionError as e:
+            logger.warning(f"Failed to load extension {extension}: {e}")
+
+Discovery
++++++++++++
+
+:meth:`.Bot.find_extensions` (and by extension, :meth:`.Bot.load_extensions`) discover modules/extensions
+similar to :func:`py:pkgutil.walk_packages`; the given root package name is resolved,
+and submodules/-packages are iterated through recursively.
+
+If a package has a ``setup`` function (similar to ``my_bot.cogs.other_complex_thing`` above),
+it won't be traversed further, i.e. ``data.py`` and ``models.py`` in the example won't be considered separate extensions.
+
+Namespace packages (see `PEP 420 <https://peps.python.org/pep-0420/>`__) are not supported (other than
+the provided root package), meaning every subpackage must have an ``__init__.py`` file.

test_bot/__main__.py

@@ -2,7 +2,6 @@
 
 import asyncio
 import logging
-import os
 import sys
 import traceback
 
@@ -51,10 +50,6 @@ async def on_ready(self):
         )
         # fmt: on
 
-    def add_cog(self, cog: commands.Cog, *, override: bool = False) -> None:
-        logger.info(f"Loading cog {cog.qualified_name}.")
-        return super().add_cog(cog, override=override)
-
     async def on_command_error(self, ctx: commands.Context, error: commands.CommandError) -> None:
         msg = f"Command `{ctx.command}` failed due to `{error}`"
         logger.error(msg, exc_info=True)
@@ -126,5 +121,9 @@ async def on_message_command_error(
 
 if __name__ == "__main__":
     bot = TestBot()
-    bot.load_extensions(os.path.join(__package__, Config.cogs_folder))
+    bot.load_extensions(
+        ".cogs",
+        package=__package__,
+        load_callback=lambda e: logger.info(f"Loaded extension {e}."),
+    )
     bot.run(Config.token)

tests/ext/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: MIT

tests/ext/commands/__init__.py

@@ -0,0 +1 @@
+# SPDX-License-Identifier: MIT

tests/ext/commands/test_common_bot_base.py

@@ -0,0 +1,45 @@
+# SPDX-License-Identifier: MIT
+
+import asyncio
+from pathlib import Path
+from typing import Iterator
+from unittest import mock
+
+import pytest
+
+from disnake.ext.commands import errors
+from disnake.ext.commands.common_bot_base import CommonBotBase
+
+from ... import helpers
+
+
+class TestExtensions:
+    @pytest.fixture()
+    def module_root(self, tmpdir: Path) -> Iterator[str]:
+        with helpers.chdir_module(tmpdir):
+            yield str(tmpdir)
+
+    @pytest.fixture()
+    def bot(self):
+        with mock.patch.object(asyncio, "get_event_loop", mock.Mock()), mock.patch.object(
+            CommonBotBase, "_fill_owners", mock.Mock()
+        ):
+            bot = CommonBotBase()
+        return bot
+
+    def test_find_path_invalid(self, bot: CommonBotBase):
+        with pytest.raises(ValueError, match=r"Paths outside the cwd are not supported"):
+            bot.find_extensions("../../etc/passwd")
+
+    def test_find(self, bot: CommonBotBase, module_root: str):
+        helpers.create_dirs(module_root, {"test_cogs": {"__init__.py": "", "admin.py": ""}})
+
+        assert bot.find_extensions("test_cogs")
+
+        with pytest.raises(errors.ExtensionError, match=r"Unable to find root module 'other_cogs'"):
+            bot.find_extensions("other_cogs")
+
+        with pytest.raises(
+            errors.ExtensionError, match=r"Module 'test_cogs.admin' is not a package"
+        ):
+            bot.find_extensions(".admin", package="test_cogs")