Source code for hyrax.verbs.verb_registry

import functools
import logging
from abc import ABC
from collections.abc import Mapping

from hyrax.context import run_context



logger = logging.getLogger(__name__)





class Verb(ABC):  # noqa: B024
    """Base class for all hyrax verbs"""

    # Verbs get to define how their parser gets added to the main parser
    # This is given in case verbs do not define any keyword args for
    # subparser.add_parser()


    add_parser_kwargs: dict[str, str] = {}


    # Subclasses declare which data_request groups they require or optionally use.
    # REQUIRED_DATA_GROUPS must all be present in the data_request config.
    # OPTIONAL_DATA_GROUPS are used when present but cause no error if absent.
    # Verbs that leave both empty skip data_request validation entirely.


    REQUIRED_DATA_GROUPS: tuple[str, ...] = ()



    OPTIONAL_DATA_GROUPS: tuple[str, ...] = ()



    cli_name = "VERB"



    description = ""




    def __init__(self, config):
        """
        .. py:method:: __init__

        Overall initialization for all verbs that saves the config
        """


        self.config = config

        self.validate_data_request()




    def __init_subclass__(cls, **kwargs):
        """Give every verb a run context for the duration of its ``run()``.

        Wrapping here rather than in each verb means a verb cannot forget to
        establish a context, and cannot leave one behind when it finishes or
        raises. Verbs still record their own results directory, with
        :func:`hyrax.context.update_context`, once they have created it.

        Only ``run`` is wrapped. Every verb's ``run_cli`` delegates to ``run``,
        so wrapping both would establish two contexts per CLI invocation.
        """
        super().__init_subclass__(**kwargs)

        run = cls.__dict__.get("run")
        if run is None or getattr(run, "_hyrax_bracketed", False):
            return

        @functools.wraps(run)
        def run_with_context(self, *args, **kwargs):
            with run_context(type(self).cli_name, config=self.config):
                return run(self, *args, **kwargs)

        # functools.wraps sets __wrapped__, so inspect.signature() still reports
        # the verb's real signature. Hyrax.__getattr__ hands verb.run straight to
        # notebook users, who rely on that for help text and completion.
        run_with_context._hyrax_bracketed = True
        cls.run = run_with_context


    @classmethod


    def information(cls):
        """Returns a string describing this verb. Includes the following:
        - Name of the verb
        - Required Data Groups
        - Optional Data Groups
        - One line description of what this verb does

        If a data group is empty then it will be printed as an empty tuple.

        Returns
        -------
        str
            <name>: Data Groups: Req. (<req1>, <req2>, ...), Opt. (<opt1>, <opt2>, ...). <Description>
        """
        info = cls.cli_name + ": "
        required = "Req. " + str(cls.REQUIRED_DATA_GROUPS)
        optional = "Opt. " + str(cls.OPTIONAL_DATA_GROUPS)
        info = info + "Data groups: " + required + ", " + optional + ". " + cls.description
        return info




    def validate_data_request(self) -> None:
        """Validate the data_request configuration for this verb's known groups.

        Reads ``data_request`` from the verb's config and verifies that every
        group listed in ``REQUIRED_DATA_GROUPS`` is present.  Verbs that define
        neither ``REQUIRED_DATA_GROUPS`` nor ``OPTIONAL_DATA_GROUPS`` skip
        validation entirely.

        Raises
        ------
        RuntimeError
            If a required group is absent from the data_request config.
        """
        if not self.REQUIRED_DATA_GROUPS and not self.OPTIONAL_DATA_GROUPS:
            return

        data_request = self.config.get("data_request")
        if not data_request:
            return

        if not isinstance(data_request, Mapping):
            raise RuntimeError(
                f"{type(self).__name__} received a non-mapping data_request configuration "
                f"of type {type(data_request)!r}; expected a mapping from group name to config."
            )

        # Verify that every required group is present in the config.
        missing = [g for g in self.REQUIRED_DATA_GROUPS if g not in data_request]
        if missing:
            raise RuntimeError(
                f"{type(self).__name__} requires dataset group(s) {missing} in the "
                f"data_request configuration, but they were not found. "
                f"Available groups: {sorted(data_request.keys())}."
            )

        # Run Pydantic structural validation so that a structurally invalid
        # data_request injected after set_config is caught at verb-instantiation
        # time rather than silently skipped.
        from pydantic import ValidationError

        from hyrax.config_schemas.data_request import DataRequestDefinition

        try:
            DataRequestDefinition.model_validate(data_request)
        except ValidationError as exc:
            raise RuntimeError(
                f"Invalid data_request configuration for {type(self).__name__}: {exc}"
            ) from exc




# Verbs with no class are assumed to have a function in hyrax.py which
# performs their function. All other verbs should be defined by named classes
# in hyrax.verbs and use the @hyrax_verb decorator


VERB_REGISTRY: dict[str, type[Verb] | None] = {
    "train": None,
    "infer": None,
    "download": None,
    "prepare": None,
    "rebuild_manifest": None,
}





def hyrax_verb(cls: type[Verb]) -> type[Verb]:
    """Decorator to register a hyrax verb"""
    from hyrax.plugin_utils import update_registry

    update_registry(VERB_REGISTRY, cls.cli_name, cls)  # type: ignore[attr-defined]
    return cls





def all_verbs() -> list[str]:
    """Returns all verbs that are currently registered"""
    return [verb for verb in VERB_REGISTRY]





def all_class_verbs() -> list[str]:
    """Returns all verbs that are currently registered with a class-based implementation"""
    return [verb for verb in VERB_REGISTRY if VERB_REGISTRY.get(verb) is not None]





def is_verb_class(cli_name: str) -> bool:
    """Returns true if the verb has a class based implementation

    Parameters
    ----------
    cli_name : str
        The name of the verb on the command line interface

    Returns
    -------
    bool
        True if the verb has a class-based implementation
    """
    return cli_name in VERB_REGISTRY and VERB_REGISTRY.get(cli_name) is not None





def fetch_verb_class(cli_name: str) -> type[Verb] | None:
    """Gives the class object for the named verb

    Parameters
    ----------
    cli_name : str
        The name of the verb on the command line interface


    Returns
    -------
    Optional[type[Verb]]
        The verb class or None if no such verb class exists.
    """
    return VERB_REGISTRY.get(cli_name)