Move .report.CALLBACKS → .report.Config.callback

- Global variable becomes an instance variable for .report.Config
  instances. This allows the use of different callbacks within the same
  program scope.
- Retain same defaults.
- Mark .report.register() function as deprecated.
- Adjust all usage of register().
- Remove fixture preserve_report_callbacks (no longer needed) and all
  usage.
- Adjust and add tests.

Author: khaeru

message_ix_models/model/transport/testing.py

@@ -176,7 +176,7 @@ def simulated_solution(request, context) -> Reporter:
     add_simulated_solution(rep, info, data)
 
     # Register the callback to set up transport reporting
-    message_ix_models.report.register(callback)
+    context.report.register(callback)
 
     # Prepare the reporter
     with silence_log("genno", logging.CRITICAL):

message_ix_models/model/transport/workflow.py

@@ -187,7 +187,7 @@ def generate(
     from message_ix_models import Workflow
     from message_ix_models.model.workflow import solve
     from message_ix_models.project.ssp import SSP_2024
-    from message_ix_models.report import register, report
+    from message_ix_models.report import report
 
     from . import build
     from .config import Config, get_cl_scenario
@@ -205,8 +205,8 @@ def generate(
     Config.from_context(context, options=options)
 
     # Set the default .report.Config key for ".* reported" steps
-    register("model.transport")
     context.report.key = report_key
+    context.report.register("model.transport")
 
     # Create the workflow
     wf = Workflow(context)

message_ix_models/report/__init__.py

@@ -1,12 +1,10 @@
 import logging
-from collections.abc import Callable
 from contextlib import nullcontext
 from copy import deepcopy
 from functools import partial
-from importlib import import_module
 from operator import itemgetter
 from pathlib import Path
-from typing import Optional, Union
+from typing import TYPE_CHECKING, Optional, Union
 from warnings import warn
 
 import genno.config
@@ -21,6 +19,9 @@
 
 from .config import Config
 
+if TYPE_CHECKING:
+    from .config import Callback
+
 __all__ = [
     "Config",
     "prepare_reporter",
@@ -43,10 +44,6 @@ def _(c: Reporter, info):
         pass
 
 
-#: List of callbacks for preparing the Reporter.
-CALLBACKS: list[Callable] = []
-
-
 @genno.config.handles("iamc")
 def iamc(c: Reporter, info):
     """Handle one entry from the ``iamc:`` config section.
@@ -119,67 +116,20 @@ def iamc(c: Reporter, info):
     c.add("concat", f"{info['variable']}::iamc", *keys)
 
 
-def register(name_or_callback: Union[Callable, str]) -> Optional[str]:
-    """Register a callback function for :meth:`prepare_reporter`.
-
-    Each registered function is called by :meth:`prepare_reporter`, in order to add or
-    modify reporting keys. Specific model variants and projects can register a callback
-    to extend the reporting graph.
-
-    Callback functions must take two arguments: the Reporter, and a :class:`.Context`:
-
-    .. code-block:: python
-
-        from message_ix.report import Reporter
-        from message_ix_models import Context
-        from message_ix_models.report import register
-
-        def cb(rep: Reporter, ctx: Context):
-            # Modify `rep` by calling its methods ...
-            pass
+def register(name_or_callback: Union["Callback", str]) -> Optional[str]:
+    """Deprecated alias for :meth:`.report.Config.register`.
 
-        register(cb)
-
-    Parameters
-    ----------
-    name_or_callback
-        If a string, this may be a submodule of :mod:`.message_ix_models`, or
-        :mod:`message_data`, in which case the function
-        ``{message_data,message_ix_models}.{name}.report.callback`` is used. Or, it may
-        be a fully-resolved package/module name, in which case ``{name}.callback`` is
-        used. If a callable (function), it is used directly.
+    This version uses :meth:`Context.get_instance()` to get the 0-th Context, and calls
+    that method.
     """
-    if isinstance(name_or_callback, str):
-        # Resolve a string
-        candidates = [
-            # As a fully-resolved package/module name
-            name_or_callback,
-            # As a submodule of message_ix_models
-            f"message_ix_models.{name_or_callback}.report",
-            # As a submodule of message_data
-            f"message_data.{name_or_callback}.report",
-        ]
-        mod = None
-        for name in candidates:
-            try:
-                mod = import_module(name)
-            except ModuleNotFoundError:
-                continue
-            else:
-                break
-        if mod is None:
-            raise ModuleNotFoundError(" or ".join(candidates))
-        callback = mod.callback
-    else:
-        callback = name_or_callback
-        name = callback.__name__
-
-    if callback in CALLBACKS:
-        log.info(f"Already registered: {callback}")
-        return None
+    warn(
+        "message_ix_models.report.register(…) function; use the method "
+        ".report.Config.register(…) or Context.report.register(…) instead",
+        DeprecationWarning,
+        stacklevel=2,
+    )
 
-    CALLBACKS.append(callback)
-    return name
+    return Context.get_instance(0).report.register(name_or_callback)
 
 
 def log_before(context, rep, key) -> None:
@@ -367,7 +317,7 @@ def prepare_reporter(
     rep.configure(model=deepcopy(context.model))
 
     # Apply callbacks for other modules which define additional reporting computations
-    for callback in CALLBACKS:
+    for callback in context.report.callback:
         callback(rep, context)
 
     key = context.report.key
@@ -415,7 +365,3 @@ def defaults(rep: Reporter, context: Context) -> None:
             rep.add(*comp, strict=True)
         except KeyExistsError:
             pass  # message_ix > 3.7.0; these are already defined
-
-
-register(defaults)
-register("message_ix_models.report.plot")

message_ix_models/report/cli.py

@@ -10,10 +10,8 @@
 
 def _modules_arg(context, param, value):
     """--module/-m: load extra reporting config from modules."""
-    from . import register
-
     for m in filter(len, value.split(",")):
-        name = register(m)
+        name = context.report.register(m)
         log.info(f"Registered reporting from {name}")
 
 

message_ix_models/report/config.py

@@ -1,16 +1,33 @@
 import logging
+from collections.abc import Callable
 from dataclasses import InitVar, dataclass, field
+from importlib import import_module
 from pathlib import Path
 from typing import TYPE_CHECKING, Optional, Union
 
 from message_ix_models.util import local_data_path, package_data_path
 from message_ix_models.util.config import ConfigHelper
 
 if TYPE_CHECKING:
+    import genno
     from genno.core.key import KeyLike
 
+    from message_ix_models.util.context import Context
+
 log = logging.getLogger(__name__)
 
+#: Type signature of callback functions referenced by :attr:`.Config.callback` and
+#: used by :func:`.prepare_reporter`.
+Callback = Callable[["genno.Computer", "Context"], None]
+
+
+def _default_callbacks() -> list[Callback]:
+    from message_ix_models.report import plot
+
+    from . import defaults
+
+    return [defaults, plot.callback]
+
 
 @dataclass
 class Config(ConfigHelper):
@@ -20,6 +37,28 @@ class Config(ConfigHelper):
     respected.
     """
 
+    #: List of callbacks for preparing the :class:`.Reporter`.
+    #:
+    #: Each registered function is called by :meth:`prepare_reporter`, in order to add
+    #: or modify reporting keys. Specific model variants and projects can register a
+    #: callback to extend the reporting graph.
+    #:
+    #: Callback functions must take two arguments: the Computer/Reporter, and a
+    #: :class:`.Context`:
+    #:
+    #: .. code-block:: python
+    #:
+    #:     from message_ix.report import Reporter
+    #:     from message_ix_models import Context
+    #:     from message_ix_models.report import register
+    #:
+    #:     def cb(rep: Reporter, ctx: Context) -> None:
+    #:         # Modify `rep` by calling its methods ...
+    #:         pass
+    #:
+    #: context.report.register(cb)
+    callback: list[Callback] = field(default_factory=_default_callbacks)
+
     #: Shorthand to call :func:`use_file` on a new instance.
     from_file: InitVar[Optional[Path]] = package_data_path("report", "global.yaml")
 
@@ -49,9 +88,54 @@ class Config(ConfigHelper):
     legacy: dict = field(default_factory=lambda: dict(use=False, merge_hist=True))
 
     def __post_init__(self, from_file, _legacy) -> None:
+        # Handle InitVars
         self.use_file(from_file)
         self.legacy.update(use=_legacy)
 
+    def register(self, name_or_callback: Union[Callback, str]) -> Optional[str]:
+        """Register a :attr:`callback` function for :func:`prepare_reporter`.
+
+        Parameters
+        ----------
+        name_or_callback
+            If a callable (function), it is used directly.
+
+            If a string, it may name a submodule of :mod:`.message_ix_models`, or
+            :mod:`message_data`, in which case the function
+            :py:`{message_data,message_ix_models}.{name}.report.callback` is used. Or,
+            it may be a fully-resolved package/module name, in which case
+            :py:`{name}.callback` is used.
+        """
+
+        if isinstance(name_or_callback, str):
+            # Resolve a string
+            candidates = [
+                name_or_callback,  # A fully-resolved package/module name
+                f"message_ix_models.{name_or_callback}.report",  # A submodule here
+                f"message_data.{name_or_callback}.report",  # A message_data submodule
+            ]
+            mod = None
+            for name in candidates:
+                try:
+                    mod = import_module(name)
+                except ModuleNotFoundError:
+                    continue
+                else:
+                    break
+            if mod is None:
+                raise ModuleNotFoundError(" or ".join(candidates))
+            callback = mod.callback
+        else:
+            callback = name_or_callback
+            name = callback.__name__
+
+        if callback in self.callback:
+            log.info(f"Already registered: {callback}")
+            return None
+
+        self.callback.append(callback)
+        return name
+
     def set_output_dir(self, arg: Optional[Path]) -> None:
         """Set :attr:`output_dir`, the output directory.
 

message_ix_models/testing/__init__.py

@@ -51,24 +51,6 @@ def pytest_addoption(parser):
 # Fixtures
 
 
-@pytest.fixture(scope="function")
-def preserve_report_callbacks():
-    """Protect :data:`.report.CALLBACKS` from effects of a test function.
-
-    Use this fixture for test functions that call :func:`.report.register` to avoid
-    changing the global/default configuration for other tests.
-    """
-
-    from message_ix_models import report
-
-    try:
-        tmp = report.CALLBACKS.copy()
-        yield
-    finally:
-        report.CALLBACKS.clear()
-        report.CALLBACKS.extend(tmp)
-
-
 @pytest.fixture(scope="session")
 def session_context(pytestconfig, tmp_env):
     """A :class:`.Context` connected to a temporary, in-memory database.
@@ -148,9 +130,6 @@ def test_context(request, session_context):
     """A copy of :func:`session_context` scoped to one test function."""
     ctx = deepcopy(session_context)
 
-    # Ensure there is a report key
-    ctx.setdefault("report", dict())
-
     yield ctx
 
     ctx.delete()

message_ix_models/tests/model/transport/test_build.py

@@ -122,7 +122,6 @@ def test_build_bare_res(
         # "ixmp://local/MESSAGEix-Transport on ENGAGE_SSP2_v4.1.7/baseline",
     ),
 )
-@pytest.mark.usefixtures("preserve_report_callbacks")
 def test_build_existing(tmp_path, test_context, url, solve=False):
     """Test that model.transport.build works on certain existing scenarios.
 

message_ix_models/tests/model/transport/test_demand.py

@@ -317,7 +317,6 @@ def test_urban_rural_shares(test_context, tmp_path, regions, years, pop_scen):
 
 @MARK[7]
 @build.get_computer.minimum_version
-@pytest.mark.usefixtures("preserve_report_callbacks")
 @pytest.mark.parametrize(
     "nodes, target",
     [

message_ix_models/tests/model/transport/test_report.py

@@ -56,7 +56,6 @@ def test_configure_legacy():
 
 @MARK[7]
 @build.get_computer.minimum_version
-@pytest.mark.usefixtures("preserve_report_callbacks")
 @pytest.mark.parametrize(
     "regions, years",
     (
@@ -75,9 +74,7 @@ def test_configure_legacy():
 def test_report_bare_solved(request, test_context, tmp_path, regions, years):
     """Run MESSAGEix-Transport–specific reporting."""
     from message_ix_models.model.transport.report import callback
-    from message_ix_models.report import Config, register
-
-    register(callback)
+    from message_ix_models.report import Config
 
     # Update configuration
     # key = "transport all"  # All including plots, etc.
@@ -87,6 +84,7 @@ def test_report_bare_solved(request, test_context, tmp_path, regions, years):
         years=years,
         report=Config("global.yaml", key=key, output_dir=tmp_path),
     )
+    test_context.report.register(callback)
 
     # Built and (optionally) solved scenario. dummy supply data is necessary for the
     # scenario to be feasible without any other contents.
@@ -114,7 +112,7 @@ def quiet_genno(caplog):
 
 @MARK[7]
 @build.get_computer.minimum_version
-@mark.usefixtures("quiet_genno", "preserve_report_callbacks")
+@mark.usefixtures("quiet_genno")
 def test_simulated_solution(request, test_context, regions="R12", years="B"):
     """:func:`message_ix_models.report.prepare_reporter` works on the simulated data."""
     test_context.update(regions=regions, years=years)
@@ -136,7 +134,7 @@ def test_simulated_solution(request, test_context, regions="R12", years="B"):
 
 @pytest.mark.xfail(condition=GHA, reason="Temporary, for #213; fails on GitHub Actions")
 @build.get_computer.minimum_version
-@mark.usefixtures("quiet_genno", "preserve_report_callbacks")
+@mark.usefixtures("quiet_genno")
 @pytest.mark.parametrize(
     "plot_name",
     # # All plots
@@ -161,7 +159,6 @@ def test_plot_simulated(request, test_context, plot_name, regions="R12", years="
 
 @pytest.mark.xfail(condition=GHA, reason="Temporary, for #213; fails on GitHub Actions")
 @sim.to_simulate.minimum_version
-@pytest.mark.usefixtures("preserve_report_callbacks")
 def test_iamc_simulated(
     request, tmp_path_factory, test_context, regions="R12", years="B"
 ) -> None: