Rework a sample test for the water module

* Use test_context as in other tests
* Type hint Context in the code
* Fix other small test-related issues

message_ix_models/model/water/__init__.py

@@ -1,3 +1,4 @@
+from .data import demands, water_supply
 from .utils import read_config
 
-__all__ = ["read_config"]
+__all__ = ["demands", "read_config", "water_supply"]

message_ix_models/model/water/build.py

@@ -1,6 +1,6 @@
 import logging
 from functools import lru_cache, partial
-from typing import Mapping
+from typing import TYPE_CHECKING, Mapping
 
 import pandas as pd
 from sdmx.model.v21 import Code
@@ -12,10 +12,13 @@
 
 from .utils import read_config
 
+if TYPE_CHECKING:
+    from message_ix_models import Context
+
 log = logging.getLogger(__name__)
 
 
-def get_spec(context) -> Mapping[str, ScenarioInfo]:
+def get_spec(context: Context) -> Mapping[str, ScenarioInfo]:
     """Return the specification for nexus implementation
 
     Parameters
@@ -190,7 +193,7 @@ def generate_set_elements(set_name, match=None):
     return results
 
 
-def map_basin(context) -> Mapping[str, ScenarioInfo]:
+def map_basin(context: Context) -> Mapping[str, ScenarioInfo]:
     """Return specification for mapping basins to regions
 
     The basins are spatially consolidated from HydroSHEDS basins delineation
@@ -242,7 +245,7 @@ def map_basin(context) -> Mapping[str, ScenarioInfo]:
     return dict(require=require, remove=remove, add=add)
 
 
-def main(context, scenario, **options):
+def main(context: Context, scenario, **options):
     """Set up MESSAGEix-Nexus on `scenario`.
 
     See also

message_ix_models/model/water/cli.py

@@ -1,10 +1,14 @@
 import logging
+from typing import TYPE_CHECKING
 
 import click
 
 from message_ix_models.model.structure import get_codes
 from message_ix_models.util.click import common_params
 
+if TYPE_CHECKING:
+    from message_ix_models import Context
+
 log = logging.getLogger(__name__)
 
 
@@ -13,12 +17,12 @@
 @common_params("regions")
 @click.option("--time", help="Manually defined time")
 @click.pass_obj
-def cli(context, regions, time):
+def cli(context: "Context", regions, time):
     """MESSAGEix-Water and Nexus variant."""
     water_ini(context, regions, time)
 
 
-def water_ini(context, regions, time):
+def water_ini(context: "Context", regions, time):
     """Add components of the MESSAGEix-Nexus module
 
     This function modifies model name & scenario name
@@ -61,7 +65,7 @@ def water_ini(context, regions, time):
     context.regions = regions
 
     # create a mapping ISO code :
-    # region name, for other scripts
+    # a region name, for other scripts
     # only needed for 1-country models
     nodes = get_codes(f"node/{context.regions}")
     nodes = list(map(str, nodes[nodes.index("World")].child))
@@ -106,7 +110,7 @@ def water_ini(context, regions, time):
     help="Defines whether the model solves with macro",
 )
 @common_params("regions")
-def nexus_cli(context, regions, rcps, sdgs, rels, macro=False):
+def nexus_cli(context: "Context", regions, rcps, sdgs, rels, macro=False):
     """
     Add basin structure connected to the energy sector and
     water balance linking different water demands to supply.
@@ -115,7 +119,7 @@ def nexus_cli(context, regions, rcps, sdgs, rels, macro=False):
     nexus(context, regions, rcps, sdgs, rels, macro)
 
 
-def nexus(context, regions, rcps, sdgs, rels, macro=False):
+def nexus(context: "Context", regions, rcps, sdgs, rels, macro=False):
     """Add basin structure connected to the energy sector and
     water balance linking different water demands to supply.
 

message_ix_models/model/water/data/__init__.py

@@ -1,6 +1,7 @@
 """Generate input data."""
 
 import logging
+from typing import TYPE_CHECKING
 
 from message_ix_models import ScenarioInfo
 from message_ix_models.util import add_par_data
@@ -11,6 +12,9 @@
 from .water_for_ppl import cool_tech, non_cooling_tec
 from .water_supply import add_e_flow, add_water_supply
 
+if TYPE_CHECKING:
+    from message_ix_models import Context
+
 log = logging.getLogger(__name__)
 
 DATA_FUNCTIONS = [
@@ -40,7 +44,7 @@
 ]
 
 
-def add_data(scenario, context, dry_run=False):
+def add_data(scenario, context: "Context", dry_run=False):
     """Populate `scenario` with MESSAGEix-Nexus data."""
 
     info = ScenarioInfo(scenario)

message_ix_models/model/water/data/demands.py

@@ -1,6 +1,7 @@
 """Prepare data for adding demands"""
 
 import os
+from typing import TYPE_CHECKING
 
 import numpy as np
 import pandas as pd
@@ -9,6 +10,9 @@
 
 from message_ix_models.util import broadcast, package_data_path
 
+if TYPE_CHECKING:
+    from message_ix_models import Context
+
 
 def get_basin_sizes(basin, node):
     """Returns the sizes of developing and developed basins for a given node"""
@@ -141,7 +145,7 @@ def target_rate_trt(df, basin):
     return df
 
 
-def add_sectoral_demands(context):
+def add_sectoral_demands(context: "Context"):
     """
     Adds water sectoral demands
     Parameters
@@ -699,7 +703,7 @@ def add_sectoral_demands(context):
     return results
 
 
-def read_water_availability(context):
+def read_water_availability(context: "Context"):
     """
     Reads water availability data and bias correct
     it for the historical years and no climate
@@ -824,7 +828,7 @@ def read_water_availability(context):
     return df_sw, df_gw
 
 
-def add_water_availability(context):
+def add_water_availability(context: "Context"):
     """
     Adds water supply constraints
     Parameters
@@ -892,7 +896,7 @@ def add_water_availability(context):
     return results
 
 
-def add_irrigation_demand(context):
+def add_irrigation_demand(context: "Context"):
     """
     Adds endogenous irrigation water demands from GLOBIOM emulator
     Parameters

message_ix_models/model/water/data/infrastructure.py

@@ -2,6 +2,7 @@
  treatment in urban & rural"""
 
 from collections import defaultdict
+from typing import TYPE_CHECKING
 
 import pandas as pd
 from message_ix import make_df
@@ -15,8 +16,11 @@
     same_time,
 )
 
+if TYPE_CHECKING:
+    from message_ix_models import Context
 
-def add_infrastructure_techs(context):  # noqa: C901
+
+def add_infrastructure_techs(context: "Context"):
     """Process water distribution data for a scenario instance.
     Parameters
     ----------
@@ -39,8 +43,7 @@ def add_infrastructure_techs(context):  # noqa: C901
     # load the scenario from context
     scen = context.get_scenario()
 
-    year_wat = [2010, 2015]
-    year_wat.extend(info.Y)
+    year_wat = (2010, 2015, *info.Y)
 
     # first activity year for all water technologies is 2020
     first_year = scen.firstmodelyear
@@ -547,7 +550,7 @@ def add_infrastructure_techs(context):  # noqa: C901
     return results
 
 
-def add_desalination(context):
+def add_desalination(context: "Context"):
     """Add desalination infrastructure
     Two types of desalination are considered;
     1. Membrane
@@ -572,8 +575,7 @@ def add_desalination(context):
     # load the scenario from context
     scen = context.get_scenario()
 
-    year_wat = [2010, 2015]
-    year_wat.extend(info.Y)
+    year_wat = (2010, 2015, *info.Y)
 
     # first activity year for all water technologies is 2020
     first_year = scen.firstmodelyear

message_ix_models/model/water/data/irrigation.py

@@ -1,13 +1,17 @@
 """Prepare data for water use for cooling & energy technologies."""
+from typing import TYPE_CHECKING
 
 import pandas as pd
 from message_ix import make_df
 
 from message_ix_models.util import broadcast, package_data_path
 
+if TYPE_CHECKING:
+    from message_ix_models import Context
+
 
 # water & electricity for irrigation
-def add_irr_structure(context):
+def add_irr_structure(context: "Context"):
     """Add irrigation withdrawal infrastructure
     The irrigation demands are added in
     Parameters

message_ix_models/model/water/data/water_for_ppl.py

@@ -1,5 +1,7 @@
 """Prepare data for water use for cooling & energy technologies."""
 
+from typing import TYPE_CHECKING, Any
+
 import numpy as np
 import pandas as pd
 from message_ix import make_df
@@ -12,9 +14,12 @@
     same_node,
 )
 
+if TYPE_CHECKING:
+    from message_ix_models import Context
+
 
 # water & electricity for cooling technologies
-def cool_tech(context):  # noqa: C901
+def cool_tech(context: "Context"):
     """Process cooling technology data for a scenario instance.
     The input values of parent technologies are read in from a scenario instance and
     then cooling fractions are calculated by using the data from
@@ -363,7 +368,7 @@ def cooling_fr(x):
     ].drop_duplicates()
     search_cols_cooling_fraction = [col for col in search_cols if col != "technology"]
 
-    def shares(x, context):
+    def shares(x, context: "Context"):
         """Process share and cooling fraction.
         Returns
         -------
@@ -379,7 +384,7 @@ def shares(x, context):
             ]["cooling_fraction"]
             x[col] = x[col] * cooling_fraction
 
-        results = []
+        results: list[Any] = []
         for i in x:
             if isinstance(i, str):
                 results.append(i)
@@ -397,7 +402,7 @@ def shares(x, context):
     hold_cost = cost[search_cols].apply(shares, axis=1, context=context)
     hold_cost = hold_cost[hold_cost["technology"] != "delme"]
 
-    def hist_act(x, context):
+    def hist_act(x, context: "Context"):
         """Calculate historical activity of cooling technology.
         The data for shares is read from ``cooltech_cost_and_shares_ssp_msg.csv``
         Returns
@@ -447,7 +452,7 @@ def hist_act(x, context):
     # dataframe for historical activities of cooling techs
     act_value_df = pd.DataFrame(changed_value_series_flat, columns=columns)
 
-    def hist_cap(x, context):
+    def hist_cap(x, context: "Context"):
         """Calculate historical capacity of cooling technology.
         The data for shares is read from ``cooltech_cost_and_shares_ssp_msg.csv``
         Returns
@@ -741,7 +746,7 @@ def hist_cap(x, context):
 
 
 # Water use & electricity for non-cooling technologies
-def non_cooling_tec(context):
+def non_cooling_tec(context: "Context"):
     """Process data for water usage of power plants (non-cooling technology related).
     Water withdrawal values for power plants are read in from
     ``tech_water_performance_ssp_msg.csv``

message_ix_models/model/water/data/water_supply.py

@@ -1,5 +1,7 @@
 """Prepare data for water use for cooling & energy technologies."""
 
+from typing import TYPE_CHECKING
+
 import numpy as np
 import pandas as pd
 from message_ix import make_df
@@ -8,8 +10,11 @@
 from message_ix_models.model.water.utils import map_yv_ya_lt
 from message_ix_models.util import broadcast, package_data_path, same_node, same_time
 
+if TYPE_CHECKING:
+    from message_ix_models import Context
+
 
-def map_basin_region_wat(context):
+def map_basin_region_wat(context: "Context"):
     """
     Calculate share of water avaialbility of basins per each parent region.
 
@@ -112,7 +117,7 @@ def map_basin_region_wat(context):
     return df_sw
 
 
-def add_water_supply(context):
+def add_water_supply(context: "Context"):
     """Add Water supply infrastructure
     This function links the water supply based on different settings and options.
     It defines the supply linkages for freshwater, groundwater and salinewater.
@@ -135,9 +140,9 @@ def add_water_supply(context):
     # load the scenario from context
     scen = context.get_scenario()
 
-    year_wat = [2010, 2015]
+    # year_wat = (2010, 2015)
     fut_year = info.Y
-    year_wat.extend(info.Y)
+    year_wat = (2010, 2015, *info.Y)
     sub_time = context.time
 
     # first activity year for all water technologies is 2020
@@ -688,7 +693,7 @@ def add_water_supply(context):
     return results
 
 
-def add_e_flow(context):
+def add_e_flow(context: "Context"):
     """Add environmental flows
     This function bounds the available water and allocates the environmental
     flows.Environmental flow bounds are calculated using Variable Monthly Flow

message_ix_models/model/water/reporting.py

@@ -260,7 +260,6 @@ def multiply_electricity_output_of_hydro(elec_hydro_var, report_iam):
 
 
 # TODO
-# flake8: noqa: C901
 def report(sc=False, reg="", sdgs=False):
     """Report nexus module results"""
 
@@ -1372,8 +1371,8 @@ def report(sc=False, reg="", sdgs=False):
                 and country_n in group["region"].values
             ):
                 report_pd.drop(group.index, inplace=True)
-                # Step 4: Rename "world" to "country" and remove rows with
-                # region = "country"
+                # Step 4: Rename "world" to "country" and remove rows
+                # with region = "country"
                 group = group[group["region"] == "World"]
                 group.loc[group["region"] == "World", "region"] = country_n
                 # Step 5: Update the original dataframe with the modified group