Fix Code Quality CI: checkout action and complex function

.github/workflows/pytest.yaml

@@ -189,7 +189,7 @@ jobs:
     steps:
     - uses: actions/checkout@v4
       with:
-        ref: ${{ needs.check.output.ref }}
+        ref: ${{ needs.check.outputs.ref }}
     - uses: astral-sh/setup-uv@v5
       with:
         cache-dependency-glob: "**/pyproject.toml"

message_ix_models/model/bare.py

@@ -115,7 +115,7 @@ def get_spec(context) -> Spec:
     nodes = get_codes(f"node/{context.model.regions}")
 
     # Top-level "World" node
-    world = nodes[nodes.index("World")]
+    world = nodes[nodes.index(Code(id="World"))]
 
     # Set elements: World, followed by the direct children of World
     add.set["node"] = [world] + world.child

message_ix_models/model/disutility.py

@@ -7,7 +7,8 @@
 
 import message_ix
 import pandas as pd
-from sdmx.model.common import Annotation, Code
+from sdmx.model.common import Code
+from sdmx.model.v21 import Annotation
 
 from message_ix_models import ScenarioInfo, Spec
 from message_ix_models.model.build import apply_spec

message_ix_models/model/material/cli.py

@@ -484,7 +484,8 @@ def test_calib(context):
     Use the --model_name and --scenario_name option to specify the scenario to solve.
     """
     # Clone and set up
-    from sdmx.model import Annotation, Code
+    from sdmx.model.common import Code
+    from sdmx.model.v21 import Annotation
 
     from message_ix_models.model import macro
     from message_ix_models.util import identify_nodes

message_ix_models/model/structure.py

@@ -11,7 +11,8 @@
 import pycountry
 import xarray as xr
 from iam_units import registry
-from sdmx.model.v21 import Annotation, Code, Codelist
+from sdmx.model.common import Code, Codelist
+from sdmx.model.v21 import Annotation
 
 from message_ix_models.util import load_package_data, package_data_path
 from message_ix_models.util.sdmx import as_codes
@@ -95,7 +96,7 @@ def get_codes(name: str) -> list[Code]:
 @cache
 def get_codelist(name: str) -> Codelist:
     """Return a :class:`.Codelist` for `name` in MESSAGEix-GLOBIOM scenarios."""
-    cl = Codelist(id=name.replace("/", "_").upper())
+    cl: Codelist = Codelist(id=name.replace("/", "_").upper())
     cl.extend(get_codes(name))
     return cl
 

message_ix_models/model/transport/build.py

@@ -395,7 +395,7 @@ def add_structure(c: Computer) -> None:
     # - `Static` tasks
     # - Single 'dynamic' tasks based on config, info, spec, and/or t_groups
     # - Multiple static and dynamic tasks generated in loops etc.
-    tasks = list(STRUCTURE_STATIC) + [
+    tasks: list[tuple] = list(STRUCTURE_STATIC) + [
         ("c::transport", quote(spec.add.set["commodity"])),
         ("c::transport+base", quote(spec.add.set["commodity"] + info.set["commodity"])),
         ("cg", quote(spec.add.set["consumer_group"])),

message_ix_models/model/water/data/water_for_ppl.py

@@ -278,6 +278,61 @@ def cooling_shares_SSP_from_yaml(
     return df_region
 
 
+def _compose_capacity_factor(inp: pd.DataFrame, context: "Context") -> pd.DataFrame:
+    """Create the capacity_factor base on data in `inp` and `context.
+
+    Parameters
+    ----------
+    inp : pd.DataFrame
+        The DataFrame representing the "input" parameter.
+    context : .Context
+        The general context in which the function is run.
+
+    Returns
+    -------
+    pd.DataFrame
+        A DataFrame representing the "capacity_factor" parameter.
+    """
+    cap_fact = make_matched_dfs(inp, capacity_factor=1)
+    # Climate Impacts on freshwater cooling capacity
+    # Taken from
+    # https://www.sciencedirect.com/science/article/pii/S0959378016301236?via%3Dihub#sec0080
+    if context.RCP == "no_climate":
+        df = cap_fact["capacity_factor"]
+    else:
+        df = cap_fact["capacity_factor"]
+        # reading ppl cooling impact dataframe
+        path = package_data_path(
+            "water", "ppl_cooling_tech", "power_plant_cooling_impact_MESSAGE.xlsx"
+        )
+        df_impact = pd.read_excel(path, sheet_name=f"{context.regions}_{context.RCP}")
+
+        for n in df_impact["node"]:
+            conditions = [
+                df["technology"].str.contains("fresh")
+                & (df["year_act"] >= 2025)
+                & (df["year_act"] < 2050)
+                & (df["node_loc"] == n),
+                df["technology"].str.contains("fresh")
+                & (df["year_act"] >= 2050)
+                & (df["year_act"] < 2070)
+                & (df["node_loc"] == n),
+                df["technology"].str.contains("fresh")
+                & (df["year_act"] >= 2070)
+                & (df["node_loc"] == n),
+            ]
+
+            choices = [
+                df_impact[(df_impact["node"] == n)]["2025s"],
+                df_impact[(df_impact["node"] == n)]["2050s"],
+                df_impact[(df_impact["node"] == n)]["2070s"],
+            ]
+
+            df["value"] = np.select(conditions, choices, default=df["value"])
+
+    return df
+
+
 # water & electricity for cooling technologies
 @minimum_version("message_ix 3.7")
 def cool_tech(context: "Context") -> dict[str, pd.DataFrame]:
@@ -855,45 +910,7 @@ def cool_tech(context: "Context") -> dict[str, pd.DataFrame]:
 
     results["technical_lifetime"] = tl
 
-    cap_fact = make_matched_dfs(inp, capacity_factor=1)
-    # Climate Impacts on freshwater cooling capacity
-    # Taken from
-    # https://www.sciencedirect.com/science/article/
-    #  pii/S0959378016301236?via%3Dihub#sec0080
-    if context.RCP == "no_climate":
-        df = cap_fact["capacity_factor"]
-    else:
-        df = cap_fact["capacity_factor"]
-        # reading ppl cooling impact dataframe
-        path = package_data_path(
-            "water", "ppl_cooling_tech", "power_plant_cooling_impact_MESSAGE.xlsx"
-        )
-        df_impact = pd.read_excel(path, sheet_name=f"{context.regions}_{context.RCP}")
-
-        for n in df_impact["node"]:
-            conditions = [
-                df["technology"].str.contains("fresh")
-                & (df["year_act"] >= 2025)
-                & (df["year_act"] < 2050)
-                & (df["node_loc"] == n),
-                df["technology"].str.contains("fresh")
-                & (df["year_act"] >= 2050)
-                & (df["year_act"] < 2070)
-                & (df["node_loc"] == n),
-                df["technology"].str.contains("fresh")
-                & (df["year_act"] >= 2070)
-                & (df["node_loc"] == n),
-            ]
-
-            choices = [
-                df_impact[(df_impact["node"] == n)]["2025s"],
-                df_impact[(df_impact["node"] == n)]["2050s"],
-                df_impact[(df_impact["node"] == n)]["2070s"],
-            ]
-
-            df["value"] = np.select(conditions, choices, default=df["value"])
-
-    results["capacity_factor"] = df
+    results["capacity_factor"] = _compose_capacity_factor(inp=inp, context=context)
 
     # Extract inand expand some paramenters from parent technologies
     # Define parameter names to be extracted

message_ix_models/project/edits/__init__.py

@@ -3,9 +3,9 @@
 from collections.abc import Callable
 from typing import TYPE_CHECKING, Any, Optional, cast
 
+import genno
 import pandas as pd
 import sdmx
-from genno import Quantity
 from genno.compat.sdmx.operator import dataset_to_quantity, quantity_to_message
 from sdmx.message import StructureMessage
 from sdmx.model.common import Codelist
@@ -57,7 +57,7 @@ def pasta_native_to_sdmx() -> "AnyQuantity":
     df = pd.read_csv(path).rename(columns={"Value": "value"}).set_index(PASTA_DIMS)
 
     # Convert to genno.Quantity
-    q = Quantity(df)
+    q = genno.Quantity(df)
 
     # Show the dimensions and codes
     print(q.coords)
@@ -208,6 +208,9 @@ def rename(qty: "AnyQuantity", name, *args, **kwargs) -> "AnyQuantity":  # type:
         },
     )
 
+    # Prevent a mypy warning in the next line; see https://github.com/khaeru/genno#166
+    assert isinstance(q_all, genno.Quantity)
+
     # Separate according to "Data" codes
     q = dict()
     keys = []

message_ix_models/project/ssp/structure.py

@@ -5,8 +5,8 @@
 from typing import TYPE_CHECKING, Optional
 
 import sdmx
-import sdmx.model.v30 as m
 import sdmx.urn
+from sdmx.model import common, v21
 
 from message_ix_models.util.sdmx import make_enum, register_agency, write
 
@@ -135,17 +135,17 @@
 def generate(context: "Context", base_dir: Optional["PathLike"] = None):
     """Generate SDMX code lists containing the SSPs."""
     # Agency for ICONICS as the maintainer of other objects
-    ICONICS = m.Agency(
+    ICONICS = common.Agency(
         id="ICONICS",
         name="International Committee on New Integrated Climate Change Assessment "
         "Scenarios",
-        contact=[m.Contact(uri=["https://depts.washington.edu/iconics/"])],
+        contact=[common.Contact(uri=["https://depts.washington.edu/iconics/"])],
     )
     register_agency(ICONICS)
 
     for cl_info in CL_INFO:
         # Create the codelist: format the name and description
-        cl: m.Codelist = m.Codelist(
+        cl: common.Codelist = common.Codelist(
             id="SSP",
             name=f"Shared Socioeconomic Pathways ({cl_info['name_extra']})",
             description="\n".join(
@@ -161,13 +161,13 @@ def generate(context: "Context", base_dir: Optional["PathLike"] = None):
             original_id = f"SSP{info['id']}"
 
             # Format the name, description; add an annotation
-            c = m.Code(
+            c = common.Code(
                 id=info["id"],
                 name=f"{original_id}: {info['name']}",
                 description="\n".join(
                     [f"{original_id}: {info['name_long']}", "", info["description"]]
                 ),
-                annotations=[m.Annotation(id="original-id", text=original_id)],
+                annotations=[v21.Annotation(id="original-id", text=original_id)],
             )
 
             # Append to the code list

message_ix_models/report/operator.py

@@ -266,6 +266,9 @@ def quantity_from_iamc(qty: "AnyQuantity", variable: str) -> "AnyQuantity":
             variables.append(match.group(0))
             replacements[match.group(0)] = match.group(1)
 
+    # Prevent a mypy warning in the next line; see https://github.com/khaeru/genno#166
+    assert isinstance(qty, genno.Quantity)
+
     subset = qty.pipe(select, {"v": variables}).pipe(relabel, {"v": replacements})
 
     unique_units = subset.coords["Unit"].data

message_ix_models/tests/model/test_disutility.py

@@ -6,7 +6,8 @@
 import pandas.testing as pdt
 import pytest
 from message_ix import make_df
-from sdmx.model.v21 import Annotation, Code
+from sdmx.model.common import Code
+from sdmx.model.v21 import Annotation
 
 from message_ix_models import ScenarioInfo, testing
 from message_ix_models.model import disutility

message_ix_models/tests/model/test_macro.py

@@ -1,7 +1,8 @@
 import pandas as pd
 import pandas.testing as pdt
 import pytest
-from sdmx.model import Annotation, Code
+from sdmx.model.common import Code
+from sdmx.model.v21 import Annotation
 
 from message_ix_models.model.macro import generate, load
 from message_ix_models.util import package_data_path

message_ix_models/tests/model/test_structure.py

@@ -5,7 +5,8 @@
 
 import pytest
 from iam_units import registry
-from sdmx.model.v21 import Annotation, Code
+from sdmx.model.common import Code
+from sdmx.model.v21 import Annotation
 
 from message_ix_models.model.structure import (
     codelists,

message_ix_models/tests/tools/costs/test_gdp.py

@@ -1,4 +1,5 @@
 import pytest
+from sdmx.model.common import Code
 
 from message_ix_models.model.structure import get_codes
 from message_ix_models.tools.costs import Config
@@ -20,7 +21,7 @@ def test_process_raw_ssp_data(test_context, node) -> None:
     # Retrieve list of node IDs
     nodes = get_codes(f"node/{node}")
     # Convert to string
-    regions = set(map(str, nodes[nodes.index("World")].child))
+    regions = set(map(str, nodes[nodes.index(Code(id="World"))].child))
 
     # Function runs
     # - context is ignored by process_raw_ssp_data
@@ -82,7 +83,7 @@ def test_adjust_cost_ratios_with_gdp(test_context, module) -> None:
     # Retrieve list of node IDs
     nodes = get_codes(f"node/{test_context.model.regions}")
     # Convert to string
-    regions = set(map(str, nodes[nodes.index("World")].child))
+    regions = set(map(str, nodes[nodes.index(Code(id="World"))].child))
 
     # Assert that all regions are present
     assert regions == set(result.region.unique())

message_ix_models/tests/util/test_sdmx.py

@@ -4,7 +4,8 @@
 
 import pytest
 import sdmx
-from sdmx.model.common import Annotation, Code
+from sdmx.model.common import Code
+from sdmx.model.v21 import Annotation
 
 from message_ix_models.util.sdmx import eval_anno, make_dataflow, make_enum, read
 

message_ix_models/tools/iamc.py

@@ -5,8 +5,8 @@
 
 import genno
 import pandas as pd
-import sdmx.model.v21 as m
 from sdmx.message import StructureMessage
+from sdmx.model import common, v21
 
 from message_ix_models.util import cached
 from message_ix_models.util.pycountry import iso_3166_alpha_3
@@ -43,8 +43,8 @@ def describe(data: pd.DataFrame, extra: Optional[str] = None) -> StructureMessag
 
     sm = StructureMessage()
 
-    def _cl(dim: str) -> m.Codelist:
-        result = m.Codelist(
+    def _cl(dim: str) -> common.Codelist:
+        result: common.Codelist = common.Codelist(
             id=dim,
             description=f"Codes appearing in the {dim!r} dimension of "
             + (extra or "data")
@@ -58,7 +58,7 @@ def _cl(dim: str) -> m.Codelist:
     for dim in ("MODEL", "SCENARIO", "REGION"):
         cl = _cl(dim)
         for value in sorted(data[dim].unique()):
-            cl.append(m.Code(id=value))
+            cl.append(common.Code(id=value))
 
     # Handle "VARIABLE" and "UNIT" jointly
     dims = ["VARIABLE", "UNIT"]
@@ -69,18 +69,18 @@ def _cl(dim: str) -> m.Codelist:
     ):
         group_units = group_data["UNIT"].unique()
         cl_variable.append(
-            m.Code(
+            common.Code(
                 id=variable,
                 annotations=[
-                    m.Annotation(
+                    v21.Annotation(
                         id="preferred-unit-measure", text=", ".join(group_units)
                     )
                 ],
             )
         )
         for unit in group_units:
             try:
-                cl_unit.append(m.Code(id=unit))
+                cl_unit.append(common.Code(id=unit))
             except ValueError:
                 pass