reloader.py

import ast
import importlib
import inspect
from os.path import samefile
from types import ModuleType
from typing import Any, Set, Tuple

import watchgod


def _try_forks(should_try: Tuple[str] = ("disnake",)) -> ModuleType:
    for module in should_try:
        try:
            return (importlib.import_module(module), module)
        except ModuleNotFoundError:
            continue


try:
    from discord.ext import commands
except ImportError:
    _, name = _try_forks()
    commands = importlib.import_module(f"{name}.ext.commands")


def _get_modules(module: ModuleType) -> Set[str]:
    visitor = _ImportVisitor()
    try:
        source = inspect.getsource(module)
    except (TypeError, OSError):
        return set()
    visitor.visit(ast.parse(source))
    return visitor.modules


def _reload_dependencies(module: ModuleType) -> Set[str]:
    modules = _get_modules(module)
    for module_name in modules:
        importlib.reload(importlib.import_module(module_name))
    return modules


class _ImportVisitor(ast.NodeVisitor):
    def __init__(self):
        self.modules: Set[str] = set()

    def visit_Import(self, node: ast.Import) -> Any:
        self.modules.update([name.name for name in node.names])

    def visit_ImportFrom(self, node: ast.ImportFrom) -> Any:
        if node.module:
            self.modules.add(node.module)

            
class ReloadCog(commands.Cog):
    def __init__(self, bot: commands.Bot) -> None:
        self.bot = bot
        self.task = self.bot.loop.create_task(self.watcher())

    def cog_unload(self):
        self.task.cancel()

    async def watcher(self) -> None:
        async for changes in watchgod.awatch("."):
            for _, file in changes:
                for name, extension in dict(self.bot.extensions).items():
                    modules = (
                        importlib.import_module(module)
                        for module in _get_modules(extension)
                    )
                    modules = (module for module in modules if module.__file__)

                    if any(
                        samefile(file, module.__file__) for module in modules  # type: ignore
                    ) or samefile(
                        extension.__file__, file  # type: ignore
                    ):

                        _reload_dependencies(extension)
                        self.bot.reload_extension(name)
                        print(f"Reloaded {name}")


def setup(bot: commands.Bot) -> None:
    bot.add_cog(ReloadCog(bot))