From 13cabfbad7826d87edc718f02c67ae3b6b18e639 Mon Sep 17 00:00:00 2001 From: Will Date: Thu, 10 Aug 2017 23:09:49 -0400 Subject: [PATCH] [Core] Add multiple/external cog paths support (#853) * WIP cog path manager * Initial working state * Get real reloading working * Type error thingy * And fix the tests * Start UI shit * path reordering * Add install path getter/setter and fix config syntax * Determine bot directory at runtime * Add UI commands for install path * Update downloader to use install path * Add sane install path default * Make evaluation of cog install path lazy * Some typing fixes * Add another line to codeowners * Conditionally put install path in paths * Always put install path first * Dont allow people to add the installdir as an additional path, guarantee install dir isn't shown with paths command * Make shit update loaded cogs * Add tests * Another one --- .github/CODEOWNERS | 1 + cogs/downloader/downloader.py | 24 +-- core/bot.py | 95 +++++++++++- core/cog_manager.py | 262 +++++++++++++++++++++++++++++++++ core/core_commands.py | 65 ++++---- core/events.py | 8 +- main.py | 12 +- tests/core/test_cog_manager.py | 57 +++++++ 8 files changed, 472 insertions(+), 52 deletions(-) create mode 100644 core/cog_manager.py create mode 100644 tests/core/test_cog_manager.py diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 25ea1b975..ca6abc505 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -3,6 +3,7 @@ # Core core/config.py @tekulvw +core/cog_manager.py @tekulvw core/drivers/* @tekulvw core/sentry_setup.py @Kowlin @tekulvw diff --git a/cogs/downloader/downloader.py b/cogs/downloader/downloader.py index 174273f6c..32d9a4d49 100644 --- a/cogs/downloader/downloader.py +++ b/cogs/downloader/downloader.py @@ -21,12 +21,6 @@ from .checks import install_agreement class Downloader: - - COG_PATH = Path.cwd() / "cogs" - LIB_PATH = Path.cwd() / "lib" - SHAREDLIB_PATH = LIB_PATH / "cog_shared" - SHAREDLIB_INIT = SHAREDLIB_PATH / "__init__.py" - def __init__(self, bot: Red): self.bot = bot @@ -40,6 +34,10 @@ class Downloader: self.already_agreed = False + self.LIB_PATH = self.bot.main_dir / "lib" + self.SHAREDLIB_PATH = self.LIB_PATH / "cog_shared" + self.SHAREDLIB_INIT = self.SHAREDLIB_PATH / "__init__.py" + self.LIB_PATH.mkdir(parents=True, exist_ok=True) self.SHAREDLIB_PATH.mkdir(parents=True, exist_ok=True) if not self.SHAREDLIB_INIT.exists(): @@ -51,6 +49,14 @@ class Downloader: self._repo_manager = RepoManager(self.conf) + @property + def cog_install_path(self): + """ + Returns the current cog install path. + :return: + """ + return self.bot.cog_mgr.install_path + @property def installed_cogs(self) -> Tuple[Installable]: """ @@ -96,7 +102,7 @@ class Downloader: """ failed = [] for cog in cogs: - if not await cog.copy_to(self.COG_PATH): + if not await cog.copy_to(self.cog_install_path): failed.append(cog) # noinspection PyTypeChecker @@ -243,7 +249,7 @@ class Downloader: " `{}`: `{}`".format(cog.name, cog.requirements)) return - await repo_name.install_cog(cog, self.COG_PATH) + await repo_name.install_cog(cog, self.cog_install_path) await self._add_to_installed(cog) @@ -260,7 +266,7 @@ class Downloader: # noinspection PyUnresolvedReferences,PyProtectedMember real_name = cog_name.name - poss_installed_path = self.COG_PATH / real_name + poss_installed_path = self.cog_install_path / real_name if poss_installed_path.exists(): await self._delete_cog(poss_installed_path) # noinspection PyTypeChecker diff --git a/core/bot.py b/core/bot.py index 01f4716d9..0c5b04bb4 100644 --- a/core/bot.py +++ b/core/bot.py @@ -1,13 +1,22 @@ +import importlib.util +from importlib.machinery import ModuleSpec + +import discord from discord.ext import commands from collections import Counter +from discord.ext.commands import GroupMixin +from pathlib import Path + from core import Config from enum import Enum import os +from core.cog_manager import CogManager + class Red(commands.Bot): - def __init__(self, cli_flags, **kwargs): + def __init__(self, cli_flags, bot_dir: Path=Path.cwd(), **kwargs): self._shutdown_mode = ExitCodes.CRITICAL self.db = Config.get_core_conf(force_registration=True) self._co_owners = cli_flags.co_owner @@ -51,6 +60,12 @@ class Red(commands.Bot): self.counter = Counter() self.uptime = None + + self.main_dir = bot_dir + + self.cog_mgr = CogManager(paths=(str(self.main_dir / 'cogs'),), + bot_dir=self.main_dir) + super().__init__(**kwargs) async def is_owner(self, user): @@ -84,12 +99,78 @@ class Red(commands.Bot): """Lists packages present in the cogs the folder""" return os.listdir("cogs") - async def save_packages_status(self): - loaded = [] - for package in self.extensions: - if package.startswith("cogs."): - loaded.append(package) - await self.db.packages.set(loaded) + async def save_packages_status(self, packages): + await self.db.packages.set(packages) + + async def add_loaded_package(self, pkg_name: str): + curr_pkgs = self.db.packages() + if pkg_name not in curr_pkgs: + curr_pkgs.append(pkg_name) + await self.save_packages_status(curr_pkgs) + + async def remove_loaded_package(self, pkg_name: str): + curr_pkgs = self.db.packages() + if pkg_name in curr_pkgs: + await self.save_packages_status([p for p in curr_pkgs if p != pkg_name]) + + def load_extension(self, spec: ModuleSpec): + name = spec.name.split('.')[-1] + if name in self.extensions: + return + + lib = spec.loader.load_module() + if not hasattr(lib, 'setup'): + del lib + raise discord.ClientException('extension does not have a setup function') + + lib.setup(self) + self.extensions[name] = lib + + def unload_extension(self, name): + lib = self.extensions.get(name) + if lib is None: + return + + lib_name = lib.__name__ # Thank you + + # find all references to the module + + # remove the cogs registered from the module + for cogname, cog in self.cogs.copy().items(): + if cog.__module__.startswith(lib_name): + self.remove_cog(cogname) + + # first remove all the commands from the module + for cmd in self.all_commands.copy().values(): + if cmd.module.startswith(lib_name): + if isinstance(cmd, GroupMixin): + cmd.recursively_remove_all_commands() + self.remove_command(cmd.name) + + # then 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__.startswith(lib_name): + remove.append(index) + + for index in reversed(remove): + del event_list[index] + + try: + func = getattr(lib, 'teardown') + except AttributeError: + pass + else: + try: + func(self) + except: + pass + finally: + # finally remove the import.. + del lib + del self.extensions[name] + # del sys.modules[name] class ExitCodes(Enum): diff --git a/core/cog_manager.py b/core/cog_manager.py new file mode 100644 index 000000000..a819313c6 --- /dev/null +++ b/core/cog_manager.py @@ -0,0 +1,262 @@ +import pkgutil +from importlib import invalidate_caches +from importlib.machinery import ModuleSpec +from typing import Tuple, Union, List +from pathlib import Path + +from discord.ext import commands + +from core import checks +from core.config import Config +from core.utils.chat_formatting import box + + +class CogManagerException(Exception): + pass + + +class InvalidPath(CogManagerException): + pass + + +class NoModuleFound(CogManagerException): + pass + + +class CogManager: + def __init__(self, paths: Tuple[str]=(), bot_dir: Path=Path.cwd()): + self.conf = Config.get_conf(self, 2938473984732, True) + self.conf.register_global( + paths=(), + install_path=str(bot_dir.resolve() / "cogs") + ) + + self._paths = set(list(self.conf.paths()) + list(paths)) + + @property + def paths(self) -> Tuple[Path, ...]: + """ + This will return all currently valid path directories. + :return: + """ + paths = [Path(p) for p in self._paths] + if self.install_path not in paths: + paths.insert(0, self.install_path) + return tuple(p.resolve() for p in paths if p.is_dir()) + + @property + def install_path(self) -> Path: + """ + Returns the install path for 3rd party cogs. + :return: + """ + p = Path(self.conf.install_path()) + return p.resolve() + + async def set_install_path(self, path: Path) -> Path: + """ + Install path setter, will return the absolute path to + the given path. + :param path: + :return: + """ + if not path.is_dir(): + raise ValueError("The install path must be an existing directory.") + resolved = path.resolve() + await self.conf.install_path.set(str(resolved)) + return resolved + + @staticmethod + def _ensure_path_obj(path: Union[Path, str]) -> Path: + """ + Guarantees an object will be a path object. + :param path: + :return: + """ + try: + path.exists() + except AttributeError: + path = Path(path) + return path + + async def add_path(self, path: Union[Path, str]): + """ + Adds a cog path to current list, will ignore duplicates. Does have + a side effect of removing all invalid paths from the saved path + list. + + Will raise InvalidPath if given anything that does not resolve to + a directory. + :param path: + :return: + """ + path = self._ensure_path_obj(path) + + # This makes the path absolute, will break if a bot install + # changes OS/Computer? + path = path.resolve() + + if not path.is_dir(): + raise InvalidPath("'{}' is not a valid directory.".format(path)) + + if path == self.install_path: + raise ValueError("Cannot add the install path as an additional path.") + + all_paths = set(self.paths + (path, )) + # noinspection PyTypeChecker + await self.set_paths(all_paths) + + async def remove_path(self, path: Union[Path, str]) -> Tuple[Path, ...]: + """ + Removes a path from the current paths list. + :param path: + :return: + """ + path = self._ensure_path_obj(path) + all_paths = list(self.paths) + if path in all_paths: + all_paths.remove(path) # Modifies in place + await self.set_paths(all_paths) + return tuple(all_paths) + + async def set_paths(self, paths_: List[Path]): + """ + Sets the current paths list. + :param paths_: + :return: + """ + self._paths = paths_ + str_paths = [str(p) for p in paths_] + await self.conf.paths.set(str_paths) + + def find_cog(self, name: str) -> ModuleSpec: + """ + Finds a cog in the list of available path. + + Raises NoModuleFound if unavailable. + :param name: + :return: + """ + resolved_paths = [str(p.resolve()) for p in self.paths] + for finder, module_name, _ in pkgutil.iter_modules(resolved_paths): + if name == module_name: + spec = finder.find_spec(name) + if spec: + return spec + + raise NoModuleFound("No module by the name of '{}' was found" + " in any available path.".format(name)) + + @staticmethod + def invalidate_caches(): + """ + This is an alias for an importlib internal and should be called + any time that a new module has been installed to a cog directory. + + *I think.* + :return: + """ + invalidate_caches() + + +class CogManagerUI: + @commands.command() + @checks.is_owner() + async def paths(self, ctx: commands.Context): + """ + Lists current cog paths in order of priority. + """ + install_path = ctx.bot.cog_mgr.install_path + cog_paths = ctx.bot.cog_mgr.paths + cog_paths = [p for p in cog_paths if p != install_path] + + msg = "Install Path: {}\n\n".format(install_path) + + partial = [] + for i, p in enumerate(cog_paths, start=1): + partial.append("{}. {}".format(i, p)) + + msg += "\n".join(partial) + await ctx.send(box(msg)) + + @commands.command() + @checks.is_owner() + async def addpath(self, ctx: commands.Context, path: Path): + """ + Add a path to the list of available cog paths. + """ + if not path.is_dir(): + await ctx.send("That path is does not exist or does not" + " point to a valid directory.") + return + + try: + await ctx.bot.cog_mgr.add_path(path) + except ValueError as e: + await ctx.send(str(e)) + else: + await ctx.send("Path successfully added.") + + @commands.command() + @checks.is_owner() + async def removepath(self, ctx: commands.Context, path_number: int): + """ + Removes a path from the available cog paths given the path_number + from !paths + """ + cog_paths = ctx.bot.cog_mgr.paths + try: + to_remove = cog_paths[path_number] + except IndexError: + await ctx.send("That is an invalid path number.") + return + + await ctx.bot.cog_mgr.remove_path(to_remove) + await ctx.send("Path successfully removed.") + + @commands.command() + @checks.is_owner() + async def reorderpath(self, ctx: commands.Context, from_: int, to: int): + """ + Reorders paths internally to allow discovery of different cogs. + """ + # Doing this because in the paths command they're 1 indexed + from_ -= 1 + to -= 1 + + all_paths = list(ctx.bot.cog_mgr.paths) + try: + to_move = all_paths.pop(from_) + except IndexError: + await ctx.send("Invalid 'from' index.") + return + + try: + all_paths.insert(to, to_move) + except IndexError: + await ctx.send("Invalid 'to' index.") + return + + await ctx.bot.cog_mgr.set_paths(all_paths) + await ctx.send("Paths reordered.") + + @commands.command() + @checks.is_owner() + async def installpath(self, ctx: commands.Context, path: Path=None): + """ + Returns the current install path or sets it if one is provided. + The provided path must be absolute or relative to the bot's + directory and it must already exist. + + No installed cogs will be transferred in the process. + """ + if path: + try: + await ctx.bot.cog_mgr.set_install_path(path) + except ValueError: + await ctx.send("That path does not exist.") + return + + install_path = ctx.bot.cog_mgr.install_path + await ctx.send("The bot will install new cogs to the `{}`" + " directory.".format(install_path)) diff --git a/core/core_commands.py b/core/core_commands.py index f549e2728..1bd667390 100644 --- a/core/core_commands.py +++ b/core/core_commands.py @@ -1,3 +1,4 @@ +import itertools from discord.ext import commands from core import checks from string import ascii_letters, digits @@ -5,11 +6,13 @@ from random import SystemRandom from collections import namedtuple import logging import importlib -import os +import sys import discord import aiohttp import asyncio +from core.cog_manager import NoModuleFound + log = logging.getLogger("red") OWNER_DISCLAIMER = ("⚠ **Only** the person who is hosting Red should be " @@ -25,29 +28,30 @@ class Core: @checks.is_owner() async def load(self, ctx, *, cog_name: str): """Loads a package""" - if not cog_name.startswith("cogs."): - cog_name = "cogs." + cog_name + try: + spec = ctx.bot.cog_mgr.find_cog(cog_name) + except NoModuleFound: + await ctx.send("No module by that name was found in any" + " cog path.") + return try: - ctx.bot.load_extension(cog_name) + ctx.bot.load_extension(spec) except Exception as e: log.exception("Package loading failed", exc_info=e) await ctx.send("Failed to load package. Check your console or " "logs for details.") else: - await ctx.bot.save_packages_status() + await ctx.bot.add_loaded_package(cog_name) await ctx.send("Done.") @commands.group() @checks.is_owner() async def unload(self, ctx, *, cog_name: str): """Unloads a package""" - if not cog_name.startswith("cogs."): - cog_name = "cogs." + cog_name - if cog_name in ctx.bot.extensions: ctx.bot.unload_extension(cog_name) - await ctx.bot.save_packages_status() + await ctx.bot.remove_loaded_package(cog_name) await ctx.send("Done.") else: await ctx.send("That extension is not loaded.") @@ -56,17 +60,11 @@ class Core: @checks.is_owner() async def _reload(self, ctx, *, cog_name: str): """Reloads a package""" - if cog_name == "downloader": - await ctx.send("DONT RELOAD DOWNLOADER.") - return - - if not cog_name.startswith("cogs."): - cog_name = "cogs." + cog_name - + ctx.bot.unload_extension(cog_name) + self.cleanup_and_refresh_modules(cog_name) try: - self.refresh_modules(cog_name) - ctx.bot.unload_extension(cog_name) - ctx.bot.load_extension(cog_name) + spec = ctx.bot.cog_mgr.find_cog(cog_name) + ctx.bot.load_extension(spec) except Exception as e: log.exception("Package reloading failed", exc_info=e) await ctx.send("Failed to reload package. Check your console or " @@ -75,18 +73,25 @@ class Core: await ctx.bot.save_packages_status() await ctx.send("Done.") - def refresh_modules(self, module): + def cleanup_and_refresh_modules(self, module_name: str): """Interally reloads modules so that changes are detected""" - module = module.replace(".", os.sep) - for root, dirs, files in os.walk(module): - for name in files: - if name.endswith(".py"): - path = os.path.join(root, name) - path, _ = os.path.splitext(path) - path = ".".join(path.split(os.sep)) - print("Reloading " + path) - m = importlib.import_module(path) - importlib.reload(m) + splitted = module_name.split('.') + + def maybe_reload(new_name): + try: + lib = sys.modules[new_name] + except KeyError: + pass + else: + importlib._bootstrap._exec(lib.__spec__, lib) + + modules = itertools.accumulate(splitted, lambda old, next: "{}.{}".format(old, next)) + for m in modules: + maybe_reload(m) + + children = {name: lib for name, lib in sys.modules.items() if name.startswith(module_name)} + for child_name, lib in children.items(): + importlib._bootstrap._exec(lib.__spec__, lib) @commands.group(name="set") async def _set(self, ctx): diff --git a/core/events.py b/core/events.py index 3ee72efcc..14ccc7dca 100644 --- a/core/events.py +++ b/core/events.py @@ -38,14 +38,12 @@ def init_events(bot, cli_flags): for package in packages: try: - bot.load_extension(package) + spec = bot.cog_mgr.find_cog(package) + bot.load_extension(spec) except Exception as e: log.exception("Failed to load package {}".format(package), exc_info=e) - failed.append(package) - - if failed: - await bot.save_packages_status() + await bot.remove_loaded_package(package) if packages: print("Loaded packages: " + ", ".join(packages)) diff --git a/main.py b/main.py index ed41e82b8..8197a7a9a 100644 --- a/main.py +++ b/main.py @@ -1,4 +1,7 @@ +from pathlib import Path + from core.bot import Red, ExitCodes +from core.cog_manager import CogManagerUI from core.global_checks import init_global_checks from core.events import init_events from core.sentry_setup import init_sentry_logging @@ -66,15 +69,22 @@ def init_loggers(cli_flags): return logger, sentry_logger +def determine_main_folder() -> Path: + return Path(os.path.dirname(__file__)).resolve() + + if __name__ == '__main__': cli_flags = parse_cli_flags() log, sentry_log = init_loggers(cli_flags) description = "Red v3 - Alpha" - red = Red(cli_flags, description=description, pm_help=None) + bot_dir = determine_main_folder() + red = Red(cli_flags, description=description, pm_help=None, + bot_dir=bot_dir) init_global_checks(red) init_events(red, cli_flags) red.add_cog(Core()) + red.add_cog(CogManagerUI()) if cli_flags.dev: red.add_cog(Dev()) diff --git a/tests/core/test_cog_manager.py b/tests/core/test_cog_manager.py new file mode 100644 index 000000000..ab4470418 --- /dev/null +++ b/tests/core/test_cog_manager.py @@ -0,0 +1,57 @@ +from pathlib import Path + +import pytest +from core import cog_manager + + +@pytest.fixture() +def cog_mgr(red): + return red.cog_mgr + + +@pytest.fixture() +def default_dir(red): + return red.main_dir + + +def test_ensure_cogs_in_paths(cog_mgr, default_dir): + cogs_dir = default_dir / 'cogs' + assert cogs_dir in cog_mgr.paths + + +@pytest.mark.asyncio +async def test_install_path_set(cog_mgr: cog_manager.CogManager, tmpdir): + path = Path(str(tmpdir)) + await cog_mgr.set_install_path(path) + assert cog_mgr.install_path == path + + +@pytest.mark.asyncio +async def test_install_path_set_bad(cog_mgr): + path = Path('something') + + with pytest.raises(ValueError): + await cog_mgr.set_install_path(path) + + +@pytest.mark.asyncio +async def test_add_path(cog_mgr, tmpdir): + path = Path(str(tmpdir)) + await cog_mgr.add_path(path) + assert path in cog_mgr.paths + + +@pytest.mark.asyncio +async def test_add_path_already_install_path(cog_mgr, tmpdir): + path = Path(str(tmpdir)) + await cog_mgr.set_install_path(path) + with pytest.raises(ValueError): + await cog_mgr.add_path(path) + + +@pytest.mark.asyncio +async def test_remove_path(cog_mgr, tmpdir): + path = Path(str(tmpdir)) + await cog_mgr.add_path(path) + await cog_mgr.remove_path(path) + assert path not in cog_mgr.paths