Source code for simplebench.reporters.reporter_manager.manager

"""ReporterManager for benchmark results.

This module contains the :class:`~.ReporterManager` class, which manages the registration,
retrieval, and unregistration of :class:`~simplebench.reporters.reporter.Reporter`
classes and instances.

This registry is global. When a :class:`~simplebench.reporters.reporter.Reporter` is registered,
it is added to the global registry, and when it is unregistered, it is removed from the
global registry.
"""
from __future__ import annotations

from argparse import ArgumentParser

from simplebench.exceptions import SimpleBenchKeyError, SimpleBenchTypeError, SimpleBenchValueError
from simplebench.reporters.choice import Choice
from simplebench.reporters.choices import Choices
from simplebench.reporters.csv.reporter import CSVReporter
from simplebench.reporters.graph.scatterplot.reporter import ScatterPlotReporter
from simplebench.reporters.json import JSONReporter
from simplebench.reporters.reporter import Reporter
from simplebench.reporters.rich_table.reporter import RichTableReporter

from .decorators.register_reporter import get_registered_reporters
from .exceptions import _ReporterManagerErrorTag

_PREDEFINED_REPORTERS: list[type[Reporter]] = [
    CSVReporter, ScatterPlotReporter, RichTableReporter, JSONReporter]
"""Container for all predefined Reporter classes.
- CSVReporter
- ScatterPlotReporter
- RichTableReporter
- JSONReporter
"""


[docs] class ReporterManager(): """Manager for all available :class:`~simplebench.reporters.reporter.Reporter` classes. This class maintains a registry of :class:`~simplebench.reporters.reporter.Reporter` instances and their associated CLI arguments. It allows for the registration, retrieval, and unregistration of :class:`~simplebench.reporters.reporter.Reporter` instances. Initially, it registers a set of built-in :class:`~simplebench.reporters.reporter.Reporter` instances: :class:`~simplebench.reporters.csv.reporter.CSVReporter`, :class:`~simplebench.reporters.graph.scatterplot.reporter.ScatterPlotReporter`, and :class:`~simplebench.reporters.rich_table.reporter.RichTableReporter`. New :class:`~simplebench.reporters.reporter.Reporter` instances can be added using the :meth:`~.register` method, and existing ones can be removed using the :meth:`~.unregister` or :meth:`~.unregister_by_name` methods. :class:`~simplebench.reporters.reporter.Reporter` instances are required to have unique names and CLI arguments in their :class:`~simplebench.reporters.choices.Choices` to avoid conflicts and the manager will raise exceptions if duplicates are detected during registration. .. code-block:: python reporter_manager = ReporterManager() my_custom_reporter = CustomReporter() reporter_manager.register(my_custom_reporter) """ def __init__(self, load_defaults: bool = True) -> None: """Initialize the ReporterManager. By default, this loads a set of predefined :class:`~simplebench.reporters.reporter.Reporter` instances and then any reporters registered via the :func:`~.register_reporter` decorator. If desired, this can be skipped by setting ``load_defaults`` to ``False``. In which case, the manager starts with an empty registry and :class:`~simplebench.reporters.reporter.Reporter` instances can be added via the :meth:`~.register` method. :param load_defaults: Whether to load the predefined reporters. Defaults to ``True``. :type load_defaults: bool """ self._registered_reporter_choices: Choices = Choices() self._registered_reporters: dict[str, Reporter] = {} if not load_defaults: return for predefined_reporter in _PREDEFINED_REPORTERS: self.register(predefined_reporter()) # type: ignore[reportCallIssue, call-arg] for registered_reporter in get_registered_reporters(): self.register(registered_reporter) @property def choices(self) -> Choices: """Return a :class:`~simplebench.reporters.choices.Choices` instance containing all registered reporter choices. :return: A :class:`~simplebench.reporters.choices.Choices` instance with all registered reporter choices. :rtype: :class:`~simplebench.reporters.choices.Choices` """ return self._registered_reporter_choices
[docs] def choice_for_arg(self, arg: str) -> Choice | None: """Return the :class:`~simplebench.reporters.choice.Choice` instance for a given CLI argument. :param arg: The CLI argument to look up. :type arg: str :return: The corresponding :class:`~simplebench.reporters.choice.Choice` instance, or ``None`` if not found. :rtype: :class:`~simplebench.reporters.choice.Choice` | None """ return self._registered_reporter_choices.get_choice_for_arg(arg)
[docs] def register(self, reporter: Reporter) -> None: """Register a new :class:`~simplebench.reporters.reporter.Reporter`. This method adds a new :class:`~simplebench.reporters.reporter.Reporter` to the registry. Unlike the predefined reporters and the :func:`~.register_reporter` decorator which both register by class, this method requires an instance of the :class:`~simplebench.reporters.reporter.Reporter`. This allows for more flexibility, such as registering :class:`~simplebench.reporters.reporter.Reporter` instances with custom initialization parameters. :param reporter: The :class:`~simplebench.reporters.reporter.Reporter` to register. :type reporter: :class:`~simplebench.reporters.reporter.Reporter` :raises SimpleBenchValueError: If a :class:`~simplebench.reporters.reporter.Reporter` with the same name or CLI argument is already registered. :raises SimpleBenchTypeError: If the provided reporter is not an instance of :class:`~simplebench.reporters.reporter.Reporter`. """ if type(reporter) is Reporter: # pylint: disable=unidiomatic-typecheck raise SimpleBenchValueError( "Cannot register the base Reporter class itself, please register a subclass instead.", tag=_ReporterManagerErrorTag.CANNOT_REGISTER_BASE_CLASS ) if not isinstance(reporter, Reporter): raise SimpleBenchTypeError( "reporter must be an instance of Reporter", tag=_ReporterManagerErrorTag.REGISTER_INVALID_REPORTER_ARG ) choices: Choices = reporter.choices if not isinstance(choices, Choices): raise SimpleBenchTypeError( "reporter.choices must return a Choices instance", tag=_ReporterManagerErrorTag.REGISTER_INVALID_CHOICES_RETURNED ) all_choice_flags: set[str] = self._registered_reporter_choices.all_choice_flags() for choice in choices.values(): if not isinstance(choice, Choice): print(f'Invalid choice type: {choice} ({type(choice)})') raise SimpleBenchTypeError(( "reporter.choices must return a Choices instance containing only Choice instances: " f'Found {choices} {type(choice)} from reporter {reporter.name!r}'), tag=_ReporterManagerErrorTag.REGISTER_INVALID_CHOICES_CONTENT ) if choice.name in self._registered_reporter_choices: raise SimpleBenchValueError( f"A reporter with the name '{choice.name}' is already registered.", tag=_ReporterManagerErrorTag.REGISTER_DUPLICATE_NAME ) for flag in choice.flags: if flag in all_choice_flags: raise SimpleBenchValueError( f"A reporter with the same CLI argument '{flag}' is already registered.", tag=_ReporterManagerErrorTag.REGISTER_DUPLICATE_CLI_ARG ) self._registered_reporter_choices.extend(choices) self._registered_reporters[reporter.name] = reporter
[docs] def all_reporters(self) -> dict[str, Reporter]: """Return all available :class:`~simplebench.reporters.reporter.Reporter` instances as a dictionary keyed by their names. :return: A dictionary of all :class:`~simplebench.reporters.reporter.Reporter` instances with their names as keys. :rtype: dict[str, :class:`~simplebench.reporters.reporter.Reporter`] """ return {reporter.name: reporter for reporter in self._registered_reporters.values()}
[docs] def unregister(self, reporter: Reporter) -> None: """Unregister a :class:`~simplebench.reporters.reporter.Reporter` by an instance exemplar. It looks up the reporter by its name and removes it from the registry. .. code-block:: python reporter_manager.unregister(MyCustomReporter()) :param reporter: :class:`~simplebench.reporters.reporter.Reporter` instance exemplar to unregister. :type reporter: :class:`~simplebench.reporters.reporter.Reporter` :raises SimpleBenchKeyError: If no reporter with the given name is registered. :raises SimpleBenchTypeError: If the provided reporter is not an instance of :class:`~simplebench.reporters.reporter.Reporter`. """ if not isinstance(reporter, Reporter): raise SimpleBenchTypeError( "reporter must be an instance of Reporter", tag=_ReporterManagerErrorTag.UNREGISTER_INVALID_REPORTER_ARG ) if reporter.name not in self._registered_reporters: raise SimpleBenchKeyError( f"No reporter with the name '{reporter.name}' is registered.", tag=_ReporterManagerErrorTag.UNREGISTER_UNKNOWN_NAME ) for choice in reporter.choices.values(): del self._registered_reporter_choices[choice.name] del self._registered_reporters[reporter.name]
[docs] def unregister_by_name(self, name: str) -> None: """Unregister a :class:`~simplebench.reporters.reporter.Reporter` by its name. :param name: The name of the :class:`~simplebench.reporters.reporter.Reporter` to unregister. :type name: str :raises SimpleBenchKeyError: If no reporter with the given name is registered. """ if name not in self._registered_reporters: raise SimpleBenchKeyError( f"No reporter with the name '{name}' is registered.", tag=_ReporterManagerErrorTag.UNREGISTER_UNKNOWN_NAME ) self.unregister(self._registered_reporters[name])
[docs] def unregister_all(self) -> None: """Unregister all :class:`~simplebench.reporters.reporter.Reporter` instances. This clears the entire registry of :class:`~simplebench.reporters.reporter.Reporter` instances. """ self._registered_reporter_choices.clear() self._registered_reporters.clear()
[docs] def add_reporters_to_argparse(self, parser: ArgumentParser) -> None: """Add all registered reporter choices to an :class:`~argparse.ArgumentParser`. This method adds the CLI arguments for all registered reporters to the provided :class:`~argparse.ArgumentParser` instance using the :meth:`~simplebench.reporters.reporter.Reporter.add_flags_to_argparse` method of each :class:`~simplebench.reporters.reporter.Reporter`. :param parser: The :class:`~argparse.ArgumentParser` to add the flags to. :type parser: :class:`~argparse.ArgumentParser` """ if not isinstance(parser, ArgumentParser): raise SimpleBenchTypeError( f'parser must be an ArgumentParser instance - cannot be a {type(parser)}', tag=_ReporterManagerErrorTag.ADD_REPORTERS_TO_ARGPARSE_INVALID_PARSER_ARG ) # Add flags for each registered reporter # The reporter is responsible for adding its own flags, so we just call its method # here. This allows each reporter to define its own flags and behavior. # We assume that the reporter's add_flags_to_argparse method is implemented correctly # and will handle any errors internally. for reporter in self._registered_reporters.values(): reporter.add_flags_to_argparse(parser=parser)