[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
This commit is contained in:
Will 2017-08-10 23:09:49 -04:00 committed by GitHub
parent 0651a6ddc3
commit 13cabfbad7
8 changed files with 472 additions and 52 deletions

1
.github/CODEOWNERS vendored
View File

@ -3,6 +3,7 @@
# Core
core/config.py @tekulvw
core/cog_manager.py @tekulvw
core/drivers/* @tekulvw
core/sentry_setup.py @Kowlin @tekulvw

View File

@ -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

View File

@ -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):

262
core/cog_manager.py Normal file
View File

@ -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))

View File

@ -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):

View File

@ -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))

12
main.py
View File

@ -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())

View File

@ -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