diff --git a/changelog/796.deprecate.rst b/changelog/796.deprecate.rst new file mode 100644 index 0000000000..1b24b620bf --- /dev/null +++ b/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. diff --git a/changelog/796.feature.rst b/changelog/796.feature.rst new file mode 100644 index 0000000000..52a9af6221 --- /dev/null +++ b/changelog/796.feature.rst @@ -0,0 +1,3 @@ +|commands| Improve :func:`Bot.load_extensions `, add :func:`Bot.find_extensions `. +- Better support for more complex extension hierarchies. +- New ``package``, ``ignore``, and ``load_callback`` parameters. diff --git a/disnake/ext/commands/common_bot_base.py b/disnake/ext/commands/common_bot_base.py index eae4cda4a4..8135377662 100644 --- a/disnake/ext/commands/common_bot_base.py +++ b/disnake/ext/commands/common_bot_base.py @@ -16,9 +16,11 @@ Callable, Dict, Generic, + Iterable, List, Mapping, Optional, + Sequence, Set, TypeVar, Union, @@ -472,7 +474,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 as e: + except (ValueError, ImportError) as e: # 3.8 raises ValueError instead of ImportError raise errors.ExtensionNotFound(name) from e def load_extension(self, name: str, *, package: Optional[str] = None) -> None: @@ -623,18 +625,149 @@ 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, + ) -> Sequence[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 + ------- + Sequence[: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 tuple(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]: diff --git a/disnake/utils.py b/disnake/utils.py index b0edf2fafa..2ef9719e6d 100644 --- a/disnake/utils.py +++ b/disnake/utils.py @@ -6,6 +6,7 @@ import asyncio import datetime import functools +import importlib import json import os import pkgutil @@ -1275,9 +1276,12 @@ def format_dt(dt: Union[datetime.datetime, float], /, style: TimestampStyle = "f return f"" +@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` @@ -1311,6 +1315,49 @@ 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 not ispkg: + yield name + continue + + # it's a package here + 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) + + def as_valid_locale(locale: str) -> Optional[str]: """Converts the provided locale name to a name that is valid for use with the API, for example by returning ``en-US`` for ``en_US``. diff --git a/docs/ext/commands/extensions.rst b/docs/ext/commands/extensions.rst index b0c178a920..fbb76d090a 100644 --- a/docs/ext/commands/extensions.rst +++ b/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 `__) are not supported (other than +the provided root package), meaning every subpackage must have an ``__init__.py`` file. diff --git a/test_bot/__main__.py b/test_bot/__main__.py index d8aef8b680..7058695edf 100644 --- a/test_bot/__main__.py +++ b/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) -> None: ) # fmt: on - def add_cog(self, cog: commands.Cog, *, override: bool = False) -> None: - logger.info("Loading cog %s", 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("Loaded extension %s.", e), + ) bot.run(Config.token) diff --git a/tests/ext/__init__.py b/tests/ext/__init__.py new file mode 100644 index 0000000000..548d2d447d --- /dev/null +++ b/tests/ext/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: MIT diff --git a/tests/ext/commands/__init__.py b/tests/ext/commands/__init__.py new file mode 100644 index 0000000000..548d2d447d --- /dev/null +++ b/tests/ext/commands/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: MIT diff --git a/tests/ext/commands/test_common_bot_base.py b/tests/ext/commands/test_common_bot_base.py new file mode 100644 index 0000000000..fc91ccc3fa --- /dev/null +++ b/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) -> None: + 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) -> None: + 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") diff --git a/tests/helpers.py b/tests/helpers.py index 2d5a4d8e41..f80b3fab22 100644 --- a/tests/helpers.py +++ b/tests/helpers.py @@ -1,10 +1,27 @@ # SPDX-License-Identifier: MIT +from __future__ import annotations + import asyncio +import contextlib import datetime import functools +import os +import sys import types -from typing import TYPE_CHECKING, Callable, ContextManager, Optional, Type, TypeVar +from pathlib import Path +from typing import ( + TYPE_CHECKING, + Any, + Callable, + ContextManager, + Dict, + Iterator, + Optional, + Type, + TypeVar, + Union, +) from unittest import mock if TYPE_CHECKING: @@ -68,3 +85,26 @@ def wrap_sync(*args, **kwargs): return func(*args, **kwargs) return wrap_sync # type: ignore + + +def create_dirs(parent: Union[str, Path], data: Dict[str, Any]) -> None: + parent = Path(parent) if isinstance(parent, str) else parent + for name, value in data.items(): + path = parent / name + if isinstance(value, dict): + path.mkdir() + create_dirs(path, value) + elif isinstance(value, str): + path.write_text(value) + + +@contextlib.contextmanager +def chdir_module(path: Union[str, Path]) -> Iterator[None]: + orig_cwd = os.getcwd() + try: + os.chdir(path) + sys.path.insert(0, str(path)) + yield + finally: + os.chdir(orig_cwd) + sys.path.remove(str(path)) diff --git a/tests/test_utils.py b/tests/test_utils.py index 1626c52fbf..879ff55e74 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -3,11 +3,11 @@ import asyncio import datetime import inspect -import os import sys import warnings from dataclasses import dataclass from datetime import timedelta, timezone +from pathlib import Path from typing import Any, Dict, List, Literal, Optional, Tuple, Union from unittest import mock @@ -808,60 +808,88 @@ def test_format_dt(dt, style, expected) -> None: @pytest.fixture(scope="session") -def tmp_module_root(tmp_path_factory): - # this obviously isn't great code, but it'll do just fine for tests +def tmp_module_root(tmp_path_factory: pytest.TempPathFactory): tmpdir = tmp_path_factory.mktemp("module_root") - for d in ["empty", "not_a_module", "mod/sub1/sub2"]: - (tmpdir / d).mkdir(parents=True) - for f in [ - "test.py", - "not_a_module/abc.py", - "mod/__init__.py", - "mod/ext.py", - "mod/sub1/sub2/__init__.py", - "mod/sub1/sub2/abc.py", - ]: - (tmpdir / f).touch() - return tmpdir + setup = "def setup(bot): ..." + helpers.create_dirs( + tmpdir, + { + "a": { + "__init__.py": "", + "nosetup.py": "", + "withsetup.py": setup, + "empty_dir": {}, + "not_a_module": {"abc.py": setup}, + "a_module": {"__init__.py": "", "abc.py": setup}, + "uncool_ext": {"__init__.py": ""}, + "cool_ext": {"__init__.py": setup}, + "mod": { + "__init__.py": "", + "ext.py": setup, + "not_a_submodule": { + "sub": {"__init__.py": setup}, + }, + "sub": { + "__init__.py": "", + "sub1": {"__init__.py": "", "abc.py": setup, "def.py": setup}, + "sub2": {"__init__.py": setup, "abc.py": setup}, + }, + }, + }, + }, + ) -@pytest.mark.parametrize( - ("path", "expected"), - [ - (".", ["test", "mod.ext"]), - ("./", ["test", "mod.ext"]), - ("empty/", []), - ], -) -def test_search_directory(tmp_module_root, path, expected) -> None: - orig_cwd = os.getcwd() - try: - os.chdir(tmp_module_root) - - # test relative and absolute paths - for p in [path, os.path.abspath(path)]: - assert sorted(utils.search_directory(p)) == sorted(expected) - finally: - os.chdir(orig_cwd) + with helpers.chdir_module(tmpdir): + yield tmpdir @pytest.mark.parametrize( - ("path", "exc"), + ("ignore", "expected"), [ - ("../../", r"Modules outside the cwd require a package to be specified"), - ("nonexistent", r"Provided path '.*?nonexistent' does not exist"), - ("test.py", r"Provided path '.*?test.py' is not a directory"), + ( + None, + [ + "a.nosetup", + "a.withsetup", + "a.a_module.abc", + "a.cool_ext", + "a.mod.ext", + "a.mod.sub.sub1.abc", + "a.mod.sub.sub1.def", + "a.mod.sub.sub2", + ], + ), + ( + ["a.nosetup", "a.mod.sub.sub1.abc", "a.mod.ext"], + [ + "a.withsetup", + "a.a_module.abc", + "a.cool_ext", + "a.mod.sub.sub1.def", + "a.mod.sub.sub2", + ], + ), + ( + lambda name: "ext" in name, # pyright: ignore[reportUnknownLambdaType] + [ + "a.nosetup", + "a.withsetup", + "a.a_module.abc", + "a.mod.sub.sub1.abc", + "a.mod.sub.sub1.def", + "a.mod.sub.sub2", + ], + ), ], ) -def test_search_directory_exc(tmp_module_root, path, exc) -> None: - orig_cwd = os.getcwd() - try: - os.chdir(tmp_module_root) - - with pytest.raises(ValueError, match=exc): - list(utils.search_directory(tmp_module_root / path)) - finally: - os.chdir(orig_cwd) +def test_walk_modules(tmp_module_root: Path, ignore, expected) -> None: + path = str(tmp_module_root / "a") + assert sorted(utils._walk_modules([path], "a.", ignore)) == sorted(expected) + + +def test_walk_modules_nonexistent(tmp_module_root: Path) -> None: + assert list(utils._walk_modules([str(tmp_module_root / "doesnotexist")], "doesnotexist.")) == [] @pytest.mark.parametrize(