Skip to content

Commit

Permalink
Update type hints for Python >= 3.9
Browse files Browse the repository at this point in the history
- Use standard collections, e.g. list[str] instead of typing.List[str].
- Import certain types from collections.abc, instead of deprecated aliases in
  typing.
- Omit .model.transport, handled in #225.
  • Loading branch information
khaeru committed Nov 20, 2024
1 parent c055458 commit c2caf20
Show file tree
Hide file tree
Showing 50 changed files with 188 additions and 188 deletions.
7 changes: 4 additions & 3 deletions message_ix_models/model/build.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import logging
from typing import Callable, Dict, List, Mapping, Optional, Union
from collections.abc import Callable, Mapping
from typing import Optional, Union

import ixmp
import pandas as pd
Expand Down Expand Up @@ -77,7 +78,7 @@ def apply_spec( # noqa: C901
pass
maybe_check_out(scenario)

dump: Dict[str, pd.DataFrame] = {} # Removed data
dump: dict[str, pd.DataFrame] = {} # Removed data

# Sort the list of sets by the number of dimensions; this places basic (non-indexed)
# sets first. Elements for these sets must be added before elements for indexed
Expand Down Expand Up @@ -166,7 +167,7 @@ def apply_spec( # noqa: C901
)


def ellipsize(elements: List) -> str:
def ellipsize(elements: list) -> str:
"""Generate a short string representation of `elements`.
If the list has more than 5 elements, only the first two and last two are shown,
Expand Down
6 changes: 3 additions & 3 deletions message_ix_models/model/disutility.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,9 @@
import logging
from collections import defaultdict
from collections.abc import Mapping, MutableMapping, Sequence
from copy import copy
from functools import partial
from itertools import product
from typing import List, Mapping, MutableMapping, Sequence

import message_ix
import pandas as pd
Expand Down Expand Up @@ -166,10 +166,10 @@ def data_conversion(info, spec: Spec) -> MutableMapping[str, pd.DataFrame]:
)

# Use the spec to retrieve information
technology: List[Code] = spec.add.set["technology"]
technology: list[Code] = spec.add.set["technology"]

# Data to return
data0: Mapping[str, List[pd.DataFrame]] = defaultdict(list)
data0: Mapping[str, list[pd.DataFrame]] = defaultdict(list)

# Loop over conversion technologies
for t in technology:
Expand Down
4 changes: 2 additions & 2 deletions message_ix_models/model/emissions.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import logging
import re
from typing import Optional, Tuple
from typing import Optional

import pandas as pd
from genno import Quantity
Expand Down Expand Up @@ -153,7 +153,7 @@ def add_tax_emission(
scen.add_par(name, data)


def split_species(unit_expr: str) -> Tuple[str, Optional[str]]:
def split_species(unit_expr: str) -> tuple[str, Optional[str]]:
"""Split `unit_expr` to an expression without a unit mention, and maybe species."""
if match := re.fullmatch("(.*)(CO2|C)(.*)", unit_expr):
return f"{match.group(1)}{match.group(3)}", match.group(2)
Expand Down
5 changes: 3 additions & 2 deletions message_ix_models/model/macro.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,10 +5,11 @@
"""

import logging
from collections.abc import Mapping
from functools import lru_cache
from itertools import product
from pathlib import Path
from typing import TYPE_CHECKING, List, Literal, Mapping, Optional, Union
from typing import TYPE_CHECKING, Literal, Optional, Union

import pandas as pd

Expand All @@ -29,7 +30,7 @@
def generate(
parameter: Literal["aeei", "config", "depr", "drate", "lotol"],
context: "Context",
commodities: Union[List[str], List["Code"]] = COMMODITY,
commodities: Union[list[str], list["Code"]] = COMMODITY,
value: Optional[float] = None,
) -> pd.DataFrame:
"""Generate uniform data for one :mod:`message_ix.macro` `parameter`.
Expand Down
5 changes: 3 additions & 2 deletions message_ix_models/model/material/build.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import logging
from typing import Any, Dict, Mapping
from collections.abc import Mapping
from typing import Any

import message_ix
import pandas as pd
Expand Down Expand Up @@ -206,7 +207,7 @@ def get_spec() -> Mapping[str, ScenarioInfo]:


def make_spec(regions: str, materials: str or None = SPEC_LIST) -> Spec:
sets: Dict[str, Any] = dict()
sets: dict[str, Any] = dict()
materials = ["common"] if not materials else materials
# Overrides specific to regional versions
tmp = dict()
Expand Down
8 changes: 4 additions & 4 deletions message_ix_models/model/material/data_aluminum.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
from collections import defaultdict
from typing import Dict, Iterable, List
from collections.abc import Iterable

import message_ix
import pandas as pd
Expand Down Expand Up @@ -69,7 +69,7 @@ def read_data_aluminum(
return data_alu, data_alu_rel, data_aluminum_ts


def gen_data_alu_ts(data: pd.DataFrame, nodes: list) -> Dict[str, pd.DataFrame]:
def gen_data_alu_ts(data: pd.DataFrame, nodes: list) -> dict[str, pd.DataFrame]:
"""
Generates time variable parameter data for aluminum sector
Parameters
Expand Down Expand Up @@ -330,7 +330,7 @@ def gen_data_alu_const(
glb_reg: str,
years: Iterable,
yv_ya: pd.DataFrame,
nodes: List[str],
nodes: list[str],
):
results = defaultdict(list)
for t in config["technology"]["add"]:
Expand Down Expand Up @@ -544,7 +544,7 @@ def gen_mock_demand_aluminum(scenario: message_ix.Scenario) -> pd.DataFrame:
return demand2020_al


def gen_data_alu_trade(scenario: message_ix.Scenario) -> Dict[str, pd.DataFrame]:
def gen_data_alu_trade(scenario: message_ix.Scenario) -> dict[str, pd.DataFrame]:
results = defaultdict(list)

data_trade = pd.read_csv(
Expand Down
20 changes: 10 additions & 10 deletions message_ix_models/model/material/data_methanol.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
from ast import literal_eval
from typing import TYPE_CHECKING, Dict, List
from typing import TYPE_CHECKING

import pandas as pd
import yaml
Expand Down Expand Up @@ -30,7 +30,7 @@
}


def gen_data_methanol(scenario: "Scenario") -> Dict[str, pd.DataFrame]:
def gen_data_methanol(scenario: "Scenario") -> dict[str, pd.DataFrame]:
"""
Generates data for methanol industry model
Expand Down Expand Up @@ -94,8 +94,8 @@ def gen_data_methanol(scenario: "Scenario") -> Dict[str, pd.DataFrame]:
def broadcast_nodes(
df_bc_node: pd.DataFrame,
df_final: pd.DataFrame,
node_cols: List[str],
node_cols_codes: Dict[str, pd.Series],
node_cols: list[str],
node_cols_codes: dict[str, pd.Series],
i: int,
) -> pd.DataFrame:
"""
Expand All @@ -105,8 +105,8 @@ def broadcast_nodes(
----------
df_bc_node: pd.DataFrame
df_final: pd.DataFrame
node_cols: List[str]
node_cols_codes: Dict[str, pd.Series]
node_cols: list[str]
node_cols_codes: dict[str, pd.Series]
i: int
"""
if len(node_cols) == 1:
Expand Down Expand Up @@ -153,17 +153,17 @@ def broadcast_nodes(

def broadcast_years(
df_bc_node: pd.DataFrame,
yr_col_out: List[str],
yr_cols_codes: Dict[str, List[str]],
yr_col_out: list[str],
yr_cols_codes: dict[str, list[str]],
col: str,
) -> pd.DataFrame:
"""
Broadcast years that were stored in pivoted row
Parameters
----------
df_bc_node: pd.DataFrame
yr_col_out: List[str]
yr_cols_codes: ict[str, List[str]]
yr_col_out: list[str]
yr_cols_codes: ict[str, list[str]]
col: str
"""
if len(yr_col_out) == 1:
Expand Down
7 changes: 3 additions & 4 deletions message_ix_models/model/material/data_petro.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,4 @@
from collections import defaultdict
from typing import List, Set

import message_ix
import pandas as pd
Expand Down Expand Up @@ -166,7 +165,7 @@ def get_demand_t1_with_income_elasticity(


def gen_data_petro_ts(
data_petro_ts: pd.DataFrame, results: dict[list], tec_ts: Set[str], nodes: List[str]
data_petro_ts: pd.DataFrame, results: dict[list], tec_ts: set[str], nodes: list[str]
) -> None:
for t in tec_ts:
common = dict(
Expand Down Expand Up @@ -239,7 +238,7 @@ def assign_input_outpt(
rg: str,
global_region: str,
common: dict,
nodes: List[str],
nodes: list[str],
) -> pd.DataFrame:
com = split[1]
lev = split[2]
Expand Down Expand Up @@ -295,7 +294,7 @@ def assign_input_outpt(
return df


def broadcast_to_regions(df: pd.DataFrame, global_region: str, nodes: List[str]):
def broadcast_to_regions(df: pd.DataFrame, global_region: str, nodes: list[str]):
if "node_loc" in df.columns:
if (
len(set(df["node_loc"])) == 1
Expand Down
8 changes: 4 additions & 4 deletions message_ix_models/model/material/data_steel.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
from collections import defaultdict
from typing import Dict, Iterable, List
from collections.abc import Iterable

import message_ix
import pandas as pd
Expand Down Expand Up @@ -119,7 +119,7 @@ def gen_mock_demand_steel(scenario: message_ix.Scenario) -> pd.DataFrame:


def gen_data_steel_ts(
data_steel_ts: pd.DataFrame, results: Dict[str, list], t: str, nodes: List[str]
data_steel_ts: pd.DataFrame, results: dict[str, list], t: str, nodes: list[str]
):
common = dict(
time="year",
Expand Down Expand Up @@ -213,11 +213,11 @@ def gen_data_steel_ts(

def get_data_steel_const(
data_steel: pd.DataFrame,
results: Dict[str, list],
results: dict[str, list],
params: Iterable,
t: str,
yv_ya: pd.DataFrame,
nodes: List[str],
nodes: list[str],
global_region: str,
):
for par in params:
Expand Down
3 changes: 2 additions & 1 deletion message_ix_models/model/material/data_util.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import os
from collections.abc import Mapping
from functools import lru_cache
from typing import TYPE_CHECKING, Literal, Mapping
from typing import TYPE_CHECKING, Literal

import ixmp
import message_ix
Expand Down
10 changes: 5 additions & 5 deletions message_ix_models/model/structure.py
Original file line number Diff line number Diff line change
@@ -1,10 +1,10 @@
import logging
import re
from collections import ChainMap
from collections.abc import Mapping, MutableMapping
from copy import copy
from functools import lru_cache
from itertools import product
from typing import Dict, List, Mapping, MutableMapping, Tuple

import click
import pandas as pd
Expand All @@ -20,7 +20,7 @@


@lru_cache()
def codelists(kind: str) -> List[str]:
def codelists(kind: str) -> list[str]:
"""Return a valid IDs for code lists of `kind`.
Parameters
Expand All @@ -32,7 +32,7 @@ def codelists(kind: str) -> List[str]:


@lru_cache()
def get_codes(name: str) -> List[Code]:
def get_codes(name: str) -> list[Code]:
"""Return codes for the dimension/set `name` in MESSAGE-GLOBIOM scenarios.
The information is read from :file:`data/{name}.yaml`, e.g.
Expand Down Expand Up @@ -101,15 +101,15 @@ def get_codelist(name: str) -> Codelist:


@lru_cache()
def get_region_codes(codelist: str) -> List[Code]:
def get_region_codes(codelist: str) -> list[Code]:
"""Return the codes that are children of "World" in the specified `codelist`."""
nodes = get_codes(f"node/{codelist}")
return nodes[nodes.index(Code(id="World"))].child


def generate_product(
data: Mapping, name: str, template: Code
) -> Tuple[List[Code], Dict[str, xr.DataArray]]:
) -> tuple[list[Code], dict[str, xr.DataArray]]:
"""Generates codes using a `template` by Cartesian product along ≥1 dimensions.
:func:`generate_set_elements` is called for each of the `dims`, and these values
Expand Down
2 changes: 1 addition & 1 deletion message_ix_models/model/water/build.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import logging
from collections.abc import Mapping
from functools import lru_cache, partial
from typing import Mapping

import pandas as pd
from sdmx.model.v21 import Code
Expand Down
3 changes: 2 additions & 1 deletion message_ix_models/model/water/data/demands.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,8 @@
"""Prepare data for adding demands"""

import os
from typing import TYPE_CHECKING, Literal, Sequence, Union
from collections.abc import Sequence
from typing import TYPE_CHECKING, Literal, Union

import numpy as np
import pandas as pd
Expand Down
6 changes: 3 additions & 3 deletions message_ix_models/model/water/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
from collections import defaultdict
from functools import lru_cache
from itertools import product
from typing import Optional, Tuple
from typing import Optional

import numpy as np
import pandas as pd
Expand Down Expand Up @@ -125,13 +125,13 @@ def func(row: pd.Series):


def map_yv_ya_lt(
periods: Tuple[int, ...], lt: int, ya: Optional[int] = None
periods: tuple[int, ...], lt: int, ya: Optional[int] = None
) -> pd.DataFrame:
"""All meaningful combinations of (vintage year, active year) given `periods`.
Parameters
----------
periods : Tuple[int, ...]
periods : tuple[int, ...]
A sequence of years.
lt : int, lifetime
ya : int, active year
Expand Down
6 changes: 3 additions & 3 deletions message_ix_models/model/workflow.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
"""Common steps for workflows."""

from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Dict, Optional
from typing import TYPE_CHECKING, Any, Optional

from message_ix import Scenario

Expand All @@ -21,7 +21,7 @@ class Config(ConfigHelper):

#: Information on an optional, other scenario from which to copy demand data in
#: :func:`solve` using :func:`transfer_demands`. Default: empty, do nothing.
demand_scenario: Dict = field(default_factory=dict)
demand_scenario: dict = field(default_factory=dict)

#: :obj:`True` to call :func:`.reserve_margin.res_marg.main` in :func:`solve`.
reserve_margin: bool = True
Expand All @@ -31,7 +31,7 @@ class Config(ConfigHelper):
#: To replicate the behaviour of the `macro_params` argument to
#: :meth:`.engage.ScenarioRunner.run`, which in turn sets the `convergence_issues`
#: argument to :meth:`.engage.ScenarioRunner.solve`, set max_adjustment to 0.1.
solve: Dict[str, Any] = field(
solve: dict[str, Any] = field(
default_factory=lambda: dict(model="MESSAGE-MACRO", max_adjustment=0.2)
)

Expand Down
Loading

0 comments on commit c2caf20

Please sign in to comment.