LogalDeveloper
9ac4a17ac8
CI / Formatting (push) Successful in 5s
CI / Linting (push) Successful in 5s
CI / Tests (Python 3.12) (push) Successful in 2m45s
CI / Tests (Python 3.13) (push) Successful in 2m53s
CI / Tests (Python 3.14) (push) Successful in 2m39s
CI / Type Checking (push) Successful in 9s
CI / Spelling (push) Successful in 5s
566 lines
22 KiB
Python
566 lines
22 KiB
Python
# Copyright 2026 Logan Fick
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""Dynamic module loading for Owlbot."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import importlib.util
|
|
import logging
|
|
import sys
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, cast
|
|
|
|
from . import builtin_modules
|
|
from .api.config import Config, ModuleConfig
|
|
from .api.context import ModuleContext
|
|
from .api.owncast_admin_client import OwncastAdminClient
|
|
from .api.owncast_client import OwncastClient
|
|
from .api.storage import ModuleStorage
|
|
from .api.templates import ModuleTemplates
|
|
from .builtin_modules import BUILTIN_MODULE_NAMES
|
|
from .registries.commands import CommandDispatcher, ModuleCommands
|
|
from .registries.events import EventDispatcher, ModuleEvents
|
|
from .registries.routes import ModuleRoutes, RouteDispatcher
|
|
|
|
if TYPE_CHECKING:
|
|
from types import ModuleType
|
|
|
|
from .api.http_client import HttpClient
|
|
from .api.lifecycle import LifecycleHandler
|
|
from .sessions import SessionManager
|
|
|
|
logger = logging.getLogger("owlbot.modules")
|
|
|
|
RESERVED_MODULE_NAMES: frozenset[str] = frozenset({"connect", "static"})
|
|
|
|
|
|
class ModuleLoadError(Exception):
|
|
"""Raised when a module fails to load."""
|
|
|
|
|
|
class ModuleLoader:
|
|
"""Discovers and loads Owlbot modules from built-in and user directories."""
|
|
|
|
def __init__(
|
|
self,
|
|
modules_dir: str | Path,
|
|
config: Config,
|
|
http_client: HttpClient,
|
|
session_manager: SessionManager,
|
|
) -> None:
|
|
"""Initialize the module loader.
|
|
|
|
:param modules_dir: Path to the user modules directory.
|
|
:param config: Configuration object for checking module enable/disable state.
|
|
:param http_client: Shared HTTP client for API clients and modules.
|
|
"""
|
|
self.modules_dir = Path(modules_dir).resolve()
|
|
self.config = config
|
|
self.http_client = http_client
|
|
self.session_manager = session_manager
|
|
self._module_contexts: dict[str, ModuleContext] = {}
|
|
self.loaded_modules: set[str] = set()
|
|
|
|
self.owncast_client = OwncastClient(
|
|
config.owncast_url,
|
|
config.owncast_access_token,
|
|
http_client,
|
|
)
|
|
|
|
# Optional admin client (only created when admin.enabled is true).
|
|
self.admin_client: OwncastAdminClient | None = None
|
|
if config.admin_enabled:
|
|
self.admin_client = OwncastAdminClient(
|
|
config.owncast_url,
|
|
config.admin_username,
|
|
config.admin_password,
|
|
http_client,
|
|
)
|
|
|
|
self.command_dispatcher = CommandDispatcher(
|
|
get_module_context=self._require_module_context,
|
|
owncast_client=self.owncast_client,
|
|
handler_timeout=config.handler_timeout,
|
|
loaded_modules=self.loaded_modules,
|
|
command_prefix=config.command_prefix,
|
|
session_manager=self.session_manager,
|
|
public_base_url=config.public_base_url,
|
|
)
|
|
self.event_dispatcher = EventDispatcher(
|
|
command_dispatch=self.command_dispatcher.dispatch,
|
|
get_module_context=self._require_module_context,
|
|
handler_timeout=config.handler_timeout,
|
|
session_manager=self.session_manager,
|
|
public_base_url=config.public_base_url,
|
|
)
|
|
self.route_dispatcher = RouteDispatcher(
|
|
get_module_context=self._require_module_context,
|
|
handler_timeout=config.handler_timeout,
|
|
session_manager=self.session_manager,
|
|
command_prefix=config.command_prefix,
|
|
)
|
|
|
|
self._core_template_dir = Path(__file__).resolve().parent / "templates"
|
|
|
|
logger.debug(
|
|
"ModuleLoader initialized (user modules directory: %s)",
|
|
self.modules_dir,
|
|
)
|
|
|
|
def get_module_context(self, module_name: str) -> ModuleContext | None:
|
|
"""Look up a module's context by name.
|
|
|
|
:param module_name: The module name.
|
|
:return: The ModuleContext if the module is loaded, None otherwise.
|
|
"""
|
|
return self._module_contexts.get(module_name)
|
|
|
|
def _require_module_context(self, module_name: str) -> ModuleContext:
|
|
"""Look up a module's context, raising if the module is not loaded.
|
|
|
|
Used by dispatchers where the module is guaranteed to be loaded.
|
|
|
|
:param module_name: The module name.
|
|
:return: The ModuleContext.
|
|
:raises RuntimeError: If the module is not loaded.
|
|
"""
|
|
ctx = self._module_contexts.get(module_name)
|
|
if ctx is None:
|
|
raise RuntimeError(f"Module '{module_name}' is not loaded")
|
|
return ctx
|
|
|
|
def discover_module_names(self) -> list[str]:
|
|
"""Discover all loadable modules from both built-in and user directories.
|
|
|
|
Built-in modules ship with the package. User modules are discovered
|
|
from the configured modules directory. If a user module has the same
|
|
name as a built-in, the user module takes precedence and a warning is
|
|
logged.
|
|
|
|
:return: Sorted list of module names.
|
|
"""
|
|
logger.debug("Discovering modules...")
|
|
user_names = self._discover_user_module_names()
|
|
|
|
reserved = user_names & RESERVED_MODULE_NAMES
|
|
for name in sorted(reserved):
|
|
logger.warning(
|
|
"User module '%s' uses a reserved name and will be skipped.", name
|
|
)
|
|
user_names -= reserved
|
|
|
|
overrides = user_names & BUILTIN_MODULE_NAMES
|
|
for name in sorted(overrides):
|
|
logger.info(
|
|
"User module '%s' found; built-in module of "
|
|
"the same name will be skipped.",
|
|
name,
|
|
)
|
|
|
|
merged = BUILTIN_MODULE_NAMES | user_names
|
|
modules = sorted(merged)
|
|
builtin_count = len(BUILTIN_MODULE_NAMES - user_names)
|
|
user_count = len(user_names)
|
|
logger.info(
|
|
"Discovered %d module(s) (%d built-in, %d user): %s",
|
|
len(modules),
|
|
builtin_count,
|
|
user_count,
|
|
", ".join(modules) if modules else "none",
|
|
)
|
|
return modules
|
|
|
|
def _discover_user_module_names(self) -> set[str]:
|
|
"""Scan the user modules directory for loadable modules.
|
|
|
|
Supports both single-file modules (``name.py``) and package modules
|
|
(``name/__init__.py``). Files and directories starting with underscore
|
|
are ignored. If both forms exist for the same name, the module is
|
|
listed once.
|
|
|
|
:return: Set of user module names.
|
|
"""
|
|
logger.debug("Scanning user modules directory: %s", self.modules_dir)
|
|
if not self.modules_dir.exists():
|
|
logger.debug("User modules directory does not exist: %s", self.modules_dir)
|
|
return set()
|
|
|
|
found: set[str] = set()
|
|
|
|
# Single-file modules (*.py).
|
|
for path in self.modules_dir.glob("*.py"):
|
|
if path.name.startswith("_"):
|
|
continue
|
|
found.add(path.stem)
|
|
|
|
# Package modules (directory with __init__.py).
|
|
for path in self.modules_dir.iterdir():
|
|
if not path.is_dir() or path.name.startswith("_"):
|
|
continue
|
|
if (path / "__init__.py").exists():
|
|
found.add(path.name)
|
|
|
|
if found:
|
|
logger.debug(
|
|
"Found %d user module(s): %s", len(found), ", ".join(sorted(found))
|
|
)
|
|
return found
|
|
|
|
async def load_all_modules(self) -> list[str]:
|
|
"""Discover and load all enabled modules using two-phase loading.
|
|
|
|
**Phase 1:** Import every module, create contexts, and register all
|
|
decorated handlers (``@on_event``, ``@on_command``, ``@on_route``).
|
|
|
|
**Phase 2:** Run ``@on_setup`` hooks for each successfully imported module.
|
|
|
|
This ordering guarantees that all decorated (static) commands are
|
|
registered before any module's ``@on_setup`` hooks attempt dynamic
|
|
registration, preventing conflicts when an ``@on_setup`` hook tries to
|
|
register a name already claimed by a decorated command.
|
|
|
|
:return: Names of successfully loaded modules.
|
|
"""
|
|
discovered = self.discover_module_names()
|
|
logger.info("Loading %d module(s)...", len(discovered))
|
|
|
|
imported: list[str] = []
|
|
loaded: list[str] = []
|
|
disabled_count = 0
|
|
failed_count = 0
|
|
|
|
# Phase 1: Import all modules and register decorated handlers.
|
|
for module_name in discovered:
|
|
try:
|
|
await self.load_module(module_name, _run_setup=False)
|
|
imported.append(module_name)
|
|
except (FileNotFoundError, ModuleLoadError) as e:
|
|
if "disabled in config" in str(e):
|
|
disabled_count += 1
|
|
else:
|
|
logger.exception("Module import failed.")
|
|
failed_count += 1
|
|
|
|
# Phase 2: Run @on_setup hooks for each imported module.
|
|
logger.debug(
|
|
"Import phase complete (%d imported). Running setup handlers...",
|
|
len(imported),
|
|
)
|
|
for module_name in imported:
|
|
try:
|
|
await self._run_module_setup(module_name)
|
|
loaded.append(module_name)
|
|
logger.info("Loaded module '%s'.", module_name)
|
|
except ModuleLoadError:
|
|
logger.exception("Module setup failed.")
|
|
failed_count += 1
|
|
|
|
parts = [f"{len(loaded)} loaded"]
|
|
if disabled_count:
|
|
parts.append(f"{disabled_count} disabled")
|
|
if failed_count:
|
|
parts.append(f"{failed_count} failed")
|
|
logger.info("Module loading complete: %s", ", ".join(parts))
|
|
return loaded
|
|
|
|
async def load_module(self, module_name: str, *, _run_setup: bool = True) -> None:
|
|
"""Load a single module by name.
|
|
|
|
If both a package and single-file form exist for the same name,
|
|
the package form is used. Checks if the module is enabled in
|
|
config before loading. After the module is executed, decorated
|
|
handlers (@on_event, @on_command) are scanned and registered.
|
|
If the module defines any ``@on_setup`` hooks, they are called
|
|
as the final step of loading (unless ``_run_setup`` is False,
|
|
in which case setup is deferred).
|
|
|
|
:param module_name: Name of the module to load (without .py extension).
|
|
:param _run_setup: Whether to run the module's @on_setup hooks.
|
|
When False, import and handler registration happen but setup
|
|
is deferred. Used internally by ``load_all_modules()`` to
|
|
implement two-phase loading.
|
|
:raises ValueError: If the module is already loaded.
|
|
:raises FileNotFoundError: If the module file cannot be found.
|
|
:raises ModuleLoadError: If the module is disabled in config,
|
|
or if the module fails to load or setup fails.
|
|
"""
|
|
if module_name in self.loaded_modules:
|
|
raise ValueError(f"Module '{module_name}' is already loaded")
|
|
|
|
if module_name in RESERVED_MODULE_NAMES:
|
|
raise ModuleLoadError(
|
|
f"Module name '{module_name}' is reserved and cannot be used"
|
|
)
|
|
|
|
logger.debug("Attempting to load module: %s", module_name)
|
|
|
|
if not self.config.is_module_enabled(module_name):
|
|
raise ModuleLoadError(f"Module '{module_name}' is disabled in config")
|
|
|
|
module_path = self._resolve_module_path(module_name)
|
|
|
|
try:
|
|
# Load the module into an isolated namespace to prevent conflicts.
|
|
# Using "owlbot_modules." prefix keeps these separate
|
|
# from normal Python packages.
|
|
spec = importlib.util.spec_from_file_location(
|
|
f"owlbot_modules.{module_name}", module_path
|
|
)
|
|
|
|
if spec is None or spec.loader is None:
|
|
raise ModuleLoadError(
|
|
f"Failed to create module spec for '{module_name}'"
|
|
)
|
|
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[f"owlbot_modules.{module_name}"] = module
|
|
|
|
logger.debug("Executing module: %s", module_name)
|
|
spec.loader.exec_module(module)
|
|
|
|
module_dir = module_path.parent
|
|
module_templates = ModuleTemplates(module_dir, self._core_template_dir)
|
|
|
|
scoped_config = ModuleConfig(self.config, module_name)
|
|
storage = ModuleStorage(self.config.storage_dir, module_name)
|
|
module_commands = ModuleCommands(self.command_dispatcher, module_name)
|
|
module_events = ModuleEvents(self.event_dispatcher, module_name)
|
|
module_routes = ModuleRoutes(
|
|
self.route_dispatcher, module_name, self.config.public_base_url
|
|
)
|
|
module_ctx = ModuleContext(
|
|
module_name=module_name,
|
|
config=scoped_config,
|
|
owncast_client=self.owncast_client,
|
|
storage=storage,
|
|
commands=module_commands,
|
|
events=module_events,
|
|
routes=module_routes,
|
|
http=self.http_client,
|
|
templates=module_templates,
|
|
admin_client=self.admin_client,
|
|
)
|
|
self._module_contexts[module_name] = module_ctx
|
|
|
|
self._register_module_handlers(module, module_name)
|
|
|
|
self.loaded_modules.add(module_name)
|
|
|
|
if _run_setup:
|
|
await self._run_module_setup(module_name)
|
|
logger.info("Loaded module '%s'.", module_name)
|
|
else:
|
|
logger.debug("Imported module '%s' (setup deferred).", module_name)
|
|
|
|
except ModuleLoadError:
|
|
raise
|
|
except Exception as e:
|
|
# Close storage if it was created before the failure.
|
|
cleanup_ctx = self._module_contexts.get(module_name)
|
|
if cleanup_ctx:
|
|
try:
|
|
# Framework lifecycle; private to module authors.
|
|
await cleanup_ctx.storage._close() # noqa: SLF001
|
|
except Exception:
|
|
logger.exception(
|
|
"Failed to close storage for module '%s' "
|
|
"during load error cleanup.",
|
|
module_name,
|
|
)
|
|
self._cleanup_module(module_name)
|
|
raise ModuleLoadError(f"Failed to load module '{module_name}': {e}") from e
|
|
|
|
async def unload_module(self, module_name: str) -> bool:
|
|
"""Unload a module, calling its teardown and cleaning up all state.
|
|
|
|
:param module_name: The module to unload.
|
|
:return: True if module was unloaded, False if not found.
|
|
"""
|
|
logger.debug("Unloading module '%s'", module_name)
|
|
if module_name not in self.loaded_modules:
|
|
logger.warning("Cannot unload '%s': not loaded", module_name)
|
|
return False
|
|
|
|
module_ctx = self._module_contexts.get(module_name)
|
|
|
|
module = sys.modules.get(f"owlbot_modules.{module_name}")
|
|
teardown_funcs = (
|
|
self._collect_lifecycle_handlers(module, "_owlbot_teardown")
|
|
if module
|
|
else []
|
|
)
|
|
if teardown_funcs and module_ctx:
|
|
for teardown_func in teardown_funcs:
|
|
try:
|
|
logger.debug("Running @on_teardown for module: %s", module_name)
|
|
await teardown_func(module_ctx)
|
|
logger.debug("Teardown completed for module: %s", module_name)
|
|
except Exception:
|
|
logger.exception("Teardown failed for module '%s'.", module_name)
|
|
else:
|
|
logger.debug("Module '%s' has no @on_teardown handlers.", module_name)
|
|
|
|
if module_ctx:
|
|
try:
|
|
# Framework lifecycle; private to module authors.
|
|
await module_ctx.storage._close() # noqa: SLF001
|
|
except Exception:
|
|
logger.exception(
|
|
"Failed to close storage for module '%s'.", module_name
|
|
)
|
|
|
|
self._cleanup_module(module_name)
|
|
logger.info("Unloaded module '%s'.", module_name)
|
|
return True
|
|
|
|
async def unload_all_modules(self) -> None:
|
|
"""Unload all modules, calling teardown and cleaning up all state.
|
|
|
|
Called during bot shutdown to allow modules to clean up resources.
|
|
"""
|
|
if not self.loaded_modules:
|
|
return
|
|
|
|
logger.info("Unloading %d module(s)...", len(self.loaded_modules))
|
|
for module_name in list(self.loaded_modules):
|
|
await self.unload_module(module_name)
|
|
logger.info("All module unload complete.")
|
|
|
|
async def _run_module_setup(self, module_name: str) -> None:
|
|
"""Run a module's ``@on_setup`` hooks if any are defined.
|
|
|
|
Each storage operation within a setup handler auto-commits
|
|
independently. If any handler fails, storage is closed and the
|
|
module is fully cleaned up.
|
|
|
|
:param module_name: Name of the module whose setup to run.
|
|
:raises ModuleLoadError: If any @on_setup handler raises an exception.
|
|
"""
|
|
module = sys.modules.get(f"owlbot_modules.{module_name}")
|
|
setup_funcs = (
|
|
self._collect_lifecycle_handlers(module, "_owlbot_setup") if module else []
|
|
)
|
|
if not setup_funcs:
|
|
logger.debug("Module '%s' has no @on_setup handlers.", module_name)
|
|
return
|
|
|
|
module_ctx = self._module_contexts[module_name]
|
|
try:
|
|
for setup_func in setup_funcs:
|
|
logger.debug("Running @on_setup for module: %s", module_name)
|
|
await setup_func(module_ctx)
|
|
logger.debug("Setup completed for module: %s", module_name)
|
|
except Exception as e:
|
|
# Framework lifecycle; private to module authors.
|
|
await module_ctx.storage._close() # noqa: SLF001
|
|
self._cleanup_module(module_name)
|
|
raise ModuleLoadError(
|
|
f"Setup failed for module '{module_name}': {e}"
|
|
) from e
|
|
|
|
def _resolve_module_path(self, module_name: str) -> Path:
|
|
"""Resolve the filesystem path for a module by name.
|
|
|
|
User modules are checked first (package form, then single-file).
|
|
If no user module is found and the name is a built-in, the built-in
|
|
path is returned.
|
|
|
|
:param module_name: The module name to resolve.
|
|
:return: The resolved path to the module file.
|
|
:raises FileNotFoundError: If neither form exists.
|
|
"""
|
|
# Check user modules directory first (allows overriding built-ins).
|
|
if self.modules_dir.exists():
|
|
package_path = self.modules_dir / module_name / "__init__.py"
|
|
if package_path.exists():
|
|
return package_path
|
|
|
|
single_path = self.modules_dir / f"{module_name}.py"
|
|
if single_path.exists():
|
|
return single_path
|
|
|
|
# Fall back to built-in modules.
|
|
if module_name in BUILTIN_MODULE_NAMES:
|
|
return self._resolve_builtin_module_path(module_name)
|
|
|
|
raise FileNotFoundError(f"Module not found: '{module_name}'")
|
|
|
|
@staticmethod
|
|
def _resolve_builtin_module_path(module_name: str) -> Path:
|
|
"""Locate a built-in module's ``__init__.py`` inside the package.
|
|
|
|
:param module_name: Name of the built-in module.
|
|
:return: Path to the module's ``__init__.py``.
|
|
"""
|
|
return (
|
|
Path(builtin_modules.__file__).resolve().parent
|
|
/ module_name
|
|
/ "__init__.py"
|
|
)
|
|
|
|
def _register_module_handlers(self, module: ModuleType, module_name: str) -> None:
|
|
"""Scan a module for decorated handlers and register them.
|
|
|
|
Delegates to each registry's ``register_from_module()`` method,
|
|
which knows how to find its own decorator markers.
|
|
|
|
:param module: The loaded Python module to scan.
|
|
:param module_name: Name of the module (for ownership tracking).
|
|
"""
|
|
self.event_dispatcher.register_from_module(module, module_name)
|
|
self.command_dispatcher.register_from_module(module, module_name)
|
|
self.route_dispatcher.register_from_module(module, module_name)
|
|
|
|
@staticmethod
|
|
def _collect_lifecycle_handlers(
|
|
module: ModuleType, marker: str
|
|
) -> list[LifecycleHandler]:
|
|
"""Collect callables from a module that have a given marker attribute.
|
|
|
|
Scans ``vars(module)`` for callable objects where ``getattr(obj, marker)``
|
|
is truthy. Used to find ``@on_setup`` (marker ``"_owlbot_setup"``) and
|
|
``@on_teardown`` (marker ``"_owlbot_teardown"``) handlers.
|
|
|
|
:param module: The loaded Python module to scan.
|
|
:param marker: The attribute name to look for (e.g. ``"_owlbot_setup"``).
|
|
:return: List of matching callables.
|
|
"""
|
|
return [
|
|
cast("LifecycleHandler", obj)
|
|
for obj in vars(module).values()
|
|
if callable(obj) and getattr(obj, marker, False)
|
|
]
|
|
|
|
def _cleanup_module(self, module_name: str) -> None:
|
|
"""Remove all state associated with a module.
|
|
|
|
Removes the top-level module entry and any submodule entries
|
|
(for package-style modules) from ``sys.modules``.
|
|
|
|
:param module_name: The module to clean up.
|
|
"""
|
|
prefix = f"owlbot_modules.{module_name}"
|
|
# Remove submodules first (e.g. owlbot_modules.custom_commands.handler),
|
|
# then the top-level entry.
|
|
for key in [k for k in sys.modules if k.startswith(f"{prefix}.")]:
|
|
sys.modules.pop(key, None)
|
|
sys.modules.pop(prefix, None)
|
|
self._module_contexts.pop(module_name, None)
|
|
self.loaded_modules.discard(module_name)
|
|
self.event_dispatcher.unregister_by_module(module_name)
|
|
self.command_dispatcher.unregister_by_module(module_name)
|
|
self.route_dispatcher.unregister_by_module(module_name)
|
|
logger.debug("Cleaned up module '%s'.", module_name)
|