Check backend in .scenario_setup functions

khaeru · khaeru · commit 4fb571c26f5d · 2025-05-09T17:28:37.000+02:00
- Replace "NOTE this assumes an IXMP4Backend" with call to
  on_ixmp4backend(); make functions no-ops if False.
- Import ixmp4 classes used only for hinting inside a TYPE_CHECKING
  block.
- Add compose_maps() for common pattern of calling compose_*_maps().
- Remove redundant conditions in .core.Scenario, .model.MESSAGE.
diff --git a/message_ix/core.py b/message_ix/core.py
@@ -885,21 +885,14 @@ def rename(self, name: str, mapping: Mapping[str, str], keep: bool = False) -> N
         maybe_commit(self, commit, f"Rename {name!r} using mapping {mapping}")
 
     def commit(self, comment: str) -> None:
-        if on_ixmp4backend(self):
-            # NB The following statements are similar to part of MESSAGE.run()
-            from message_ix.util.scenario_setup import (
-                compose_dimension_map,
-                compose_period_map,
-            )
-            # JDBCBackend calls these functions as part of every commit, but they have
-            # moved to message_ix because they handle message-specific data
+        from message_ix.util.scenario_setup import compose_maps
+
+        # JDBCBackend calls these functions as part of every commit, but they have moved
+        # to message_ix because they handle message-specific data
 
-            # The sanity checks fail for some tests (e.g. 'node' being empty)
-            # ensure_required_indexsets_have_data(scenario=self)
+        # The sanity checks fail for some tests (e.g. 'node' being empty)
+        # ensure_required_indexsets_have_data(scenario=self)
 
-            # Compose some auxiliary tables
-            for dimension in ("node", "time"):
-                compose_dimension_map(scenario=self, dimension=dimension)
-            compose_period_map(scenario=self)
+        compose_maps(self)
 
         return super().commit(comment)
diff --git a/message_ix/models.py b/message_ix/models.py
@@ -334,6 +334,8 @@ def initialize(cls, scenario: "ixmp.core.scenario.Scenario") -> None:
         --------
         :attr:`items`
         """
+        from message_ix.util.scenario_setup import add_default_data
+
         # Check for storage items that may contain incompatible data or need to be
         # re-initialized
         state = None
@@ -355,81 +357,70 @@ def initialize(cls, scenario: "ixmp.core.scenario.Scenario") -> None:
         # Initialize the ixmp items for MESSAGE
         cls.initialize_items(scenario, items)
 
-        if on_ixmp4backend(scenario):
-            from message_ix.util.scenario_setup import add_default_data
-
-            assert isinstance(scenario, Scenario)
-
-            if not scenario.platform._units_to_warn_about:
-                scenario.platform._units_to_warn_about = REQUIRED_UNITS.copy()
-
-            # NOTE I tried transcribing this from ixmp_source as-is, but the MESSAGE
-            # class defined in models.py takes care of setting up the Scenario -- except
-            # for adding default data.
-            # ixmp_source does other things, too, which I don't think we need here, but
-            # I've kept them in for completeness for now.
-
-            # ixmp_source first sets up a Scenario and adds default data
-            # models.MESSAGE seems to do the setup for us in all cases, while
-            # add_default_data() only adds missing items, so can always run.
-            # TODO Is this correct?
-            # if version == "new":
-            #     # If the Scenario already exists, we don't need these two
-            # set_up_scenario(s=self)
-            add_default_data(scenario=scenario)
-
-            # TODO We don't seem to need this, but if we do, give them better names
-            # self.tecParList = [
-            #     parameter_info for parameter_info in PARAMETERS if parameter_info.is_tec # noqa: E501
-            # ]
-            # self.tecActParList = [
-            #     parameter_info
-            #     for parameter_info in PARAMETERS
-            #     if parameter_info.is_tec_act
-            # ]
-
-            # TODO the following could be activated in ixmp_source through the flag
-            # parameter `sanity_checks`. This 'sanity_check' (there are more, s.b.) is
-            # generally only active when loading a scenario from the DB (unless
-            # explicitly loading via ID, in which case it's also inactive). We don't
-            # distinguish loading from the DB and some tutorials failed, so disable.
-            # ensure_required_indexsets_have_data(s=self)
-
-            # TODO It does not seem useful to construct these because some required
-            # indexsets won't have any data in them yet. They do get run in imxp_source
-            # at this point, though.
-            # compose_dimension_map(s=self, dimension="node")
-            # compose_dimension_map(s=self, dimension="time")
-            # compose_period_map(s=self)
+        assert isinstance(scenario, Scenario)
+
+        if not scenario.platform._units_to_warn_about:
+            scenario.platform._units_to_warn_about = REQUIRED_UNITS.copy()
+
+        # NOTE I tried transcribing this from ixmp_source as-is, but the MESSAGE class
+        # defined in models.py takes care of setting up the Scenario -- except for
+        # adding default data.
+        # ixmp_source does other things, too, which I don't think we need here, but I've
+        # kept them in for completeness for now.
+
+        # ixmp_source first sets up a Scenario and adds default data
+        # models.MESSAGE seems to do the setup for us in all cases, while
+        # add_default_data() only adds missing items, so can always run.
+        # TODO Is this correct?
+        # if version == "new":
+        #     # If the Scenario already exists, we don't need these two
+        # set_up_scenario(s=self)
+        add_default_data(scenario=scenario)
+
+        # TODO We don't seem to need this, but if we do, give them better names
+        # self.tecParList = [
+        #     parameter_info for parameter_info in PARAMETERS if parameter_info.is_tec
+        # ]
+        # self.tecActParList = [
+        #     parameter_info
+        #     for parameter_info in PARAMETERS
+        #     if parameter_info.is_tec_act
+        # ]
+
+        # TODO the following could be activated in ixmp_source through the flag
+        # parameter `sanity_checks`. This 'sanity_check' (there are more, s.b.) is
+        # generally only active when loading a scenario from the DB (unless explicitly
+        # loading via ID, in which case it's also inactive). We don't distinguish
+        # loading from the DB and some tutorials failed, so disable.
+        # ensure_required_indexsets_have_data(s=self)
+
+        # TODO It does not seem useful to construct these because some required
+        # indexsets won't have any data in them yet. They do get run in imxp_source at
+        # this point, though.
+        # compose_maps(scenario=scenario)
 
         # Commit if anything was removed
         maybe_commit(scenario, state, f"{cls.__name__}.initialize")
 
     def run(self, scenario: "ixmp.core.scenario.Scenario") -> None:
-        if on_ixmp4backend(scenario):
-            from message_ix.util.gams_io import (
-                add_auxiliary_items_to_container_data_list,
-                add_default_data_to_container_data_list,
-                store_message_version,
-            )
-            from message_ix.util.scenario_setup import (
-                compose_dimension_map,
-                compose_period_map,
-                ensure_required_indexsets_have_data,
-            )
-
-            assert isinstance(scenario, Scenario)  # Narrow type
+        from message_ix.util.gams_io import (
+            add_auxiliary_items_to_container_data_list,
+            add_default_data_to_container_data_list,
+            store_message_version,
+        )
+        from message_ix.util.scenario_setup import (
+            compose_maps,
+            ensure_required_indexsets_have_data,
+        )
 
-            # Run the sanity checks
-            ensure_required_indexsets_have_data(scenario=scenario)
+        assert isinstance(scenario, Scenario)  # Narrow type
 
-            # NB The following are similar to statements in .core.Scenario.commit()
-            # Compose some auxiliary tables
-            for dimension in ("node", "time"):
-                compose_dimension_map(scenario=scenario, dimension=dimension)
+        # Run the sanity checks
+        ensure_required_indexsets_have_data(scenario=scenario)
 
-            compose_period_map(scenario=scenario)
+        compose_maps(scenario=scenario)
 
+        if on_ixmp4backend(scenario):
             # ixmp.model.gams.GAMSModel.__init__() creates the container_data attribute
             # from its .defaults and any user kwargs
 
diff --git a/message_ix/util/scenario_setup.py b/message_ix/util/scenario_setup.py
@@ -4,11 +4,14 @@
 import pandas as pd
 
 if TYPE_CHECKING:
+    from ixmp4 import Run
+    from ixmp4.core import IndexSet, Parameter, Table
+
     from message_ix.core import Scenario
 
 from ixmp import Platform
-from ixmp4 import Run
-from ixmp4.core import IndexSet, Parameter, Table
+
+from message_ix.util.ixmp4 import on_ixmp4backend
 
 from .scenario_data import (
     DEFAULT_INDEXSET_DATA,
@@ -28,9 +31,11 @@
 
 #     (According to ixmp_source.)
 #     """
-#     # NOTE this assumes an IXMP4Backend
+#     if not on_ixmp4backend(scenario):
+#         return
+#
 #     # Get the Run associated with the Scenario
-#     run = cast(Run, scenario.platform._backend.index[scenario])
+#     run = cast("Run", scenario.platform._backend.index[scenario])
 
 #     # Add all required IndexSets
 #     for indexset_name in REQUIRED_INDEXSETS:
@@ -71,9 +76,11 @@
 
 def add_default_data(scenario: "Scenario") -> None:
     """Add default data expected in a MESSAGEix Scenario."""
-    # NOTE this assumes an IXMP4Backend
+    if not on_ixmp4backend(scenario):
+        return
+
     # Get the Run associated with the Scenario
-    run = cast(Run, scenario.platform._backend.index[scenario])
+    run = cast("Run", scenario.platform._backend.index[scenario])
 
     # Add IndexSet data
     for indexset_data_info in DEFAULT_INDEXSET_DATA:
@@ -121,11 +128,13 @@ def ensure_required_indexsets_have_data(scenario: "Scenario") -> None:
     ValueError
         If the required IndexSets are empty.
     """
+    if not on_ixmp4backend(scenario):
+        return
+
     indexsets_to_check = ("node", "technology", "year", "time")
 
-    # NOTE this assumes an IXMP4Backend
     # Get the Run associated with the Scenario
-    run = cast(Run, scenario.platform._backend.index[scenario])
+    run = cast("Run", scenario.platform._backend.index[scenario])
 
     # Raise an error if any of the checked IndexSets are empty
     for name in indexsets_to_check:
@@ -136,7 +145,7 @@ def ensure_required_indexsets_have_data(scenario: "Scenario") -> None:
 
 
 def _maybe_add_to_table(
-    table: Table, data: Union[dict[str, Any], pd.DataFrame]
+    table: "Table", data: Union[dict[str, Any], pd.DataFrame]
 ) -> None:
     """Add (parts of) `data` to `table` if they are missing."""
     # NOTE This function doesn't handle empty data as internally, this won't happen
@@ -166,9 +175,11 @@ def compose_dimension_map(
     dimension: 'node' or 'time'
         Whether to handle the spatial or temporal dimension.
     """
-    # NOTE this assumes an IXMP4Backend
+    if not on_ixmp4backend(scenario):
+        return
+
     # Get the Run associated with the Scenario
-    run = cast(Run, scenario.platform._backend.index[scenario])
+    run = cast("Run", scenario.platform._backend.index[scenario])
 
     # Handle both spatial and temporal dimensions
     name_part = "spatial" if dimension == "node" else "temporal"
@@ -236,15 +247,15 @@ def _find_all_descendants(parent: T) -> list[T]:
 
 
 def _maybe_add_single_item_to_indexset(
-    indexset: IndexSet, data: Union[float, int, str]
+    indexset: "IndexSet", data: Union[float, int, str]
 ) -> None:
     """Add `data` to `indexset` if it is missing."""
     if data not in list(indexset.data):
         indexset.add(data=data)
 
 
 def _maybe_add_list_to_indexset(
-    indexset: IndexSet, data: Union[list[float], list[int], list[str]]
+    indexset: "IndexSet", data: Union[list[float], list[int], list[str]]
 ) -> None:
     """Add missing parts of `data` to `indexset`."""
     # NOTE missing will always only have one type, but how to tell mypy?
@@ -255,7 +266,8 @@ def _maybe_add_list_to_indexset(
 
 
 def _maybe_add_to_indexset(
-    indexset: IndexSet, data: Union[float, int, str, list[float], list[int], list[str]]
+    indexset: "IndexSet",
+    data: Union[float, int, str, list[float], list[int], list[str]],
 ) -> None:
     """Add (parts of) `data` to `indexset` if they are missing."""
     # NOTE This function doesn't handle empty data as internally, this won't happen
@@ -268,7 +280,7 @@ def _maybe_add_to_indexset(
 # NOTE this could be combined with `_maybe_add_to_table()`, but that function would be
 # slower than necessary (though likely not by much). Is the maintenance effort worth it?
 def _maybe_add_to_parameter(
-    parameter: Parameter, data: Union[dict[str, Any], pd.DataFrame]
+    parameter: "Parameter", data: Union[dict[str, Any], pd.DataFrame]
 ) -> None:
     """Add (parts of) `data` to `parameter` if they are missing."""
     # NOTE This function doesn't handle empty data as internally, this won't happen
@@ -291,14 +303,33 @@ def _maybe_add_to_parameter(
     parameter.add(data=new_data)
 
 
+def compose_maps(scenario: "Scenario") -> None:
+    """Compose maps.
+
+    - Call :func:`compose_dimension_map` for:
+
+      - :py:`dimension="node"`
+      - :py:`dimension="time"`
+
+    - Call :func:`compose_period_map`.
+    """
+    # Compose some auxiliary tables
+    for dimension in ("node", "time"):
+        compose_dimension_map(scenario=scenario, dimension=dimension)
+
+    compose_period_map(scenario=scenario)
+
+
 def compose_period_map(scenario: "Scenario") -> None:
     """Add data to the 'duration_period' Parameter in `scenario`.
 
     This covers `assignPeriodMaps()` from ixmp_source.
     """
-    # NOTE this assumes an IXMP4Backend
+    if not on_ixmp4backend(scenario):
+        return
+
     # Get the Run associated with the Scenario
-    run = cast(Run, scenario.platform._backend.index[scenario])
+    run = cast("Run", scenario.platform._backend.index[scenario])
 
     # TODO Included here in ixmp_source; this should likely move to add_default_data
     # Add one default item to 'type_year'
