diff --git a/src/management_commands/core.py b/src/management_commands/core.py index 86e0b05..12f2ef6 100644 --- a/src/management_commands/core.py +++ b/src/management_commands/core.py @@ -1,5 +1,9 @@ from __future__ import annotations +import contextlib +import importlib +import pkgutil +import typing from contextlib import suppress from django.apps.registry import apps @@ -14,6 +18,9 @@ CommandTypeError, ) +if typing.TYPE_CHECKING: + from collections.abc import Iterator + def import_command_class(dotted_path: str) -> type[BaseCommand]: try: @@ -27,6 +34,51 @@ def import_command_class(dotted_path: str) -> type[BaseCommand]: return command_class +def iterate_modules(dotted_path: str) -> Iterator[str]: + for _, name, is_pkg in pkgutil.iter_modules( + importlib.import_module(dotted_path).__path__, + ): + if not is_pkg and not name.startswith("_"): + yield name + + +def _discover_commands_in_module(module: str) -> list[str]: + commands: list[str] = [] + try: + files_in_dir = list(iterate_modules(module)) + except ImportError: # module doesn't exist + return commands + + for file in files_in_dir: + with ( + contextlib.suppress(CommandImportError), + contextlib.suppress(CommandTypeError), + ): + import_command_class(f"{module}.{file}.Command") + commands.append(file) + + return commands + + +def get_commands_from_modules_and_submodules() -> dict[str, list[str]]: + commands = {} + for module in settings.MODULES: + if module_commands := _discover_commands_in_module(module): + commands[module] = module_commands + + for app in apps.get_app_configs(): + for submodule in settings.SUBMODULES: + if app.name == "django.core" or submodule == "management.commands": + continue + + if module_commands := _discover_commands_in_module( + f"{app.name}.{submodule}", + ): + commands[app.name] = module_commands + + return commands + + def get_command_paths(name: str, app_label: str | None = None) -> list[str]: if not app_label: app_names = [ diff --git a/src/management_commands/management.py b/src/management_commands/management.py index c046422..0a10a60 100644 --- a/src/management_commands/management.py +++ b/src/management_commands/management.py @@ -1,5 +1,6 @@ from __future__ import annotations +import itertools import sys from typing import TYPE_CHECKING @@ -7,7 +8,11 @@ from django.core.management.color import color_style from .conf import settings -from .core import import_command_class, load_command_class +from .core import ( + get_commands_from_modules_and_submodules, + import_command_class, + load_command_class, +) if TYPE_CHECKING: from django.core.management.base import BaseCommand @@ -43,11 +48,21 @@ def main_help_text(self, commands_only: bool = False) -> str: if (aliases := settings.ALIASES) else [] ) + modules = get_commands_from_modules_and_submodules() + modules_usage = ( + [ + style.NOTICE(f"[django-management-commands: {module}]"), + *[f" {file}" for file in modules[module]], + "", + ] + for module in modules + ) usage_list = usage.split("\n") usage_list.append("") usage_list.extend(commands_usage) usage_list.extend(aliases_usage) + usage_list.extend(itertools.chain(*modules_usage)) return "\n".join(usage_list) diff --git a/tests/test_core.py b/tests/test_core.py index e588057..522a444 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -8,6 +8,7 @@ from management_commands.core import ( get_command_paths, + get_commands_from_modules_and_submodules, import_command_class, load_command_class, ) @@ -119,6 +120,81 @@ def test_get_command_paths_returns_list_of_all_dotted_paths_to_command_classes_i ] +def test_get_commands_from_modules_and_submodules_returns_dictionary_of_available_commands( + mocker: MockerFixture, +) -> None: + # Configure. + mocker.patch.multiple( + "management_commands.conf.settings", + MODULES=[ + "module_a", + "module_b", # no commands + ], + SUBMODULES=[ + "submodule_a", + "submodule_b", # no commands + ], + ) + + # Arrange. + app_config_a_mock = mocker.Mock() + app_config_a_mock.name = "app_a" + app_config_b_mock = mocker.Mock() # no commands + app_config_b_mock.name = "app_b" + + class CommandA: + pass + + class CommandB(BaseCommand): + pass + + # Mock. + mocker.patch( + "management_commands.core.apps.app_configs", + { + "app_a": app_config_a_mock, + "app_b": app_config_b_mock, + }, + ) + + def import_string_side_effect(dotted_path: str) -> type: + if dotted_path == "module_a.command_a.Command": + return CommandA + if dotted_path == "module_a.command_b.Command": + return CommandB + if dotted_path == "app_a.submodule_a.command_a.Command": + return CommandA + if dotted_path == "app_a.submodule_a.command_b.Command": + return CommandB + + raise ImportError + + mocker.patch( + "management_commands.core.import_string", + side_effect=import_string_side_effect, + ) + + def iterate_modules_side_effect(dotted_path: str) -> list[str]: + if dotted_path == "module_a": + return ["command_a", "command_b"] + if dotted_path == "app_a.submodule_a": + return ["command_a", "command_b"] + raise ImportError + + mocker.patch( + "management_commands.core.iterate_modules", + side_effect=iterate_modules_side_effect, + ) + + # Act. + commands = get_commands_from_modules_and_submodules() + + # Assert. + assert set(commands.keys()) == {"module_a", "app_a"} + assert commands["module_a"] == ["command_b"] + assert commands["app_a"] == ["command_b"] + + def test_get_command_paths_returns_list_of_dotted_paths_to_app_submodules_if_app_label_specified( mocker: MockerFixture, ) -> None: diff --git a/tests/test_management.py b/tests/test_management.py index 0ab3849..270ba92 100644 --- a/tests/test_management.py +++ b/tests/test_management.py @@ -56,6 +56,71 @@ def test_execute_from_command_line_help_displays_paths_and_aliases( ) in captured.out +def test_execute_from_command_line_help_displays_modules_and_submodules( + mocker: MockerFixture, + capsys: pytest.CaptureFixture[str], +) -> None: + # Mock. + mocker.patch.multiple( + "management_commands.management.settings", + MODULES=["module_a"], + SUBMODULES=["submodule_a"], + ) + + # Arrange. + app_config_a_mock = mocker.Mock() + app_config_a_mock.name = "app_a" + + class CommandB(BaseCommand): + pass + + # Mock. + mocker.patch( + "management_commands.core.apps.app_configs", + {"app_a": app_config_a_mock}, + ) + + def import_string_side_effect(dotted_path: str) -> type: + if dotted_path == "module_a.command_b.Command": + return CommandB + if dotted_path == "app_a.submodule_a.command_b.Command": + return CommandB + + raise ImportError + + mocker.patch( + "management_commands.core.import_string", + side_effect=import_string_side_effect, + ) + + def iterate_modules_side_effect(dotted_path: str) -> list[str]: + if dotted_path == "module_a": + return ["command_b"] + if dotted_path == "app_a.submodule_a": + return ["command_b"] + raise ImportError + + mocker.patch( + "management_commands.core.iterate_modules", + side_effect=iterate_modules_side_effect, + ) + + # Act. + + execute_from_command_line(["manage.py", "--help"]) + captured = capsys.readouterr() + + # Assert. + assert ( + "[django-management-commands: module_a]\n" + " command_b\n" + "\n" + "[django-management-commands: app_a]\n" + " command_b\n" + "\n" + ) in captured.out + + def test_execute_from_command_line_falls_back_to_django_management_utility_if_command_name_is_not_passed( mocker: MockerFixture, ) -> None: