Source code for autojob.plugins

"""This module loads and stores all registered plugins.

In particular, this module loads task, analysis, and harvest plugins.
"""

from importlib.metadata import EntryPoint
from importlib.metadata import entry_points
import logging
from typing import TYPE_CHECKING
from typing import Any
from typing import Literal

if TYPE_CHECKING:
    from autojob.bases.analysis_base import AnalysisBase
    from autojob.bases.harvester_base import HarvesterBase
    from autojob.bases.task_base import TaskBase

logger = logging.getLogger(__name__)

task_plugins: dict[str, EntryPoint] = {}
analysis_plugins: dict[str, EntryPoint] = {}
harvester_plugins: dict[str, EntryPoint] = {}




def discover_task_plugins() -> None:
    """Discover tasks registered through entry points."""
    tasks = entry_points(group="autojob.task")
    task_plugins.update({ep.name: ep for ep in tasks})





def discover_analysis_plugins() -> None:
    """Discover analyses registered through entry points."""
    analyses = entry_points(group="autojob.analysis")
    analysis_plugins.update({ep.name: ep for ep in analyses})





def discover_harvester_plugins() -> None:
    """Discover harvesters registered through entry points."""
    harvesters = entry_points(group="autojob.harvest")
    harvester_plugins.update({ep.name: ep for ep in harvesters})





def get_task_class(task_class: str) -> "type[TaskBase]":
    """Return an installed task class.

    Raises:
        KeyError: The requested task class is not installed.
    """
    try:
        plugin = task_plugins[task_class]
        if isinstance(plugin, EntryPoint):
            return plugin.load()
        return plugin
    except KeyError as err:
        msg = f"Task {task_class} is not installed"
        raise ValueError(msg) from err





def get_analysis(analysis: str) -> "AnalysisBase":
    """Return an installed analysis.

    Raises:
        KeyError: The requested analysis is not installed.
    """
    try:
        plugin = analysis_plugins[analysis]
        if isinstance(plugin, EntryPoint):
            return plugin.load()
        return plugin
    except KeyError as err:
        msg = f"Analysis {analysis} is not installed"
        raise ValueError(msg) from err





def get_harvester(harvester: str) -> "HarvesterBase":
    """Return an installed harvester.

    Raises:
        KeyError: The requested harvester is not installed.
    """
    try:
        plugin = harvester_plugins[harvester]
        if isinstance(plugin, EntryPoint):
            return plugin.load()
        return plugin
    except KeyError as err:
        msg = f"Harvester {harvester} is not installed"
        raise ValueError(msg) from err





def register_plugin(
    name: str,
    plugin: Any,
    plugin_type: Literal["analysis", "harvester", "task_class"],
) -> None:
    """Register a plugin.

    Args:
        name: The name used to store the plugin.
        plugin: The plugin.
        plugin_type: The type of plugin to be registered.

    Raises:
        ValueError: Unsupported plugin type.
        TypeError: `plugin` is not the correct type.
    """
    match plugin_type:
        case "analysis":
            from autojob.bases.analysis_base import AnalysisBase  # noqa: I001, PLC0415

            type_to_verify = AnalysisBase
            verifier = isinstance
            dest = analysis_plugins
        case "harvester":
            from autojob.bases.harvester_base import HarvesterBase  # noqa: I001, PLC0415

            type_to_verify = HarvesterBase
            verifier = isinstance
            dest = harvester_plugins
        case "task_class":
            from autojob.bases.task_base import TaskBase  # noqa: PLC0415

            type_to_verify = TaskBase
            verifier = issubclass
            dest = task_plugins
        case _:
            msg = f"Unsupported plugin type {plugin_type}"
            raise ValueError(msg)

    if verifier(plugin, type_to_verify):
        dest[name] = plugin
        logger.info("Successfully registered %s plugin", name)
    else:
        msg = (
            f"{plugin_type.capitalize()} plugin must be of type "
            f"{type_to_verify.__class__.__name__}"
        )
        raise TypeError(msg)



discover_task_plugins()
discover_analysis_plugins()
discover_harvester_plugins()