Merge pull request #963 from iiasa/enh/ixmp-strict-mypy

Adjust for strict type hinting in ixmp

Author: khaeru

.gitignore

@@ -43,6 +43,7 @@ tests/data/nightly
 # Python tooling, pytest, etc.
 .benchmarks
 .coverage*
+.dmypy*
 .mypy_cache
 .pytest_cache
 .ruff_cache

.pre-commit-config.yaml

@@ -1,6 +1,6 @@
 repos:
 - repo: https://github.com/pre-commit/mirrors-mypy
-  rev: v1.14.1
+  rev: v1.17.0
   hooks:
   - id: mypy
     pass_filenames: false
@@ -15,9 +15,9 @@ repos:
     - types-PyYAML
     - types-requests
 - repo: https://github.com/astral-sh/ruff-pre-commit
-  rev: v0.9.1
+  rev: v0.12.4
   hooks:
-  - id: ruff
+  - id: ruff-check
   - id: ruff-format
     args: [ --check ]
 

RELEASE_NOTES.rst

@@ -5,6 +5,7 @@ All changes
 -----------
 
 - Document the :ref:`minimum version of Java <install-java>` required for :class:`ixmp.JDBCBackend <ixmp.backend.jdbc.JDBCBackend>` (:pull:`962`).
+- Improve type hinting (:pull:`963`).
 
 .. _v3.11.1:
 

doc/requirements.in

@@ -1,7 +1,8 @@
-# Input file for pip-compile 
+# Input file for uv pip compile
 # Only specify the packages necessary for [docs]
 
 gitpython==3.1.44
+ixmp @ git+https://github.com/iiasa/ixmp.git@main
 numpydoc==1.8.0
 pandas==2.2.3
 sphinx==8.2.3

doc/requirements.txt

@@ -1,5 +1,5 @@
 # This file was autogenerated by uv via the following command:
-#    uv pip compile --strip-extras doc/requirements.in -o doc/requirements.txt
+#    uv pip compile --strip-extras requirements.in -o requirements.txt
 alabaster==0.7.16
     # via sphinx
 babel==2.16.0
@@ -8,34 +8,85 @@ certifi==2024.7.4
     # via requests
 charset-normalizer==3.2.0
     # via requests
+click==8.2.1
+    # via
+    #   dask
+    #   ixmp
+cloudpickle==3.1.1
+    # via dask
+dask==2025.7.0
+    # via genno
 docutils==0.20.1
     # via
     #   pybtex-docutils
     #   sphinx
     #   sphinx-rtd-theme
     #   sphinxcontrib-bibtex
+et-xmlfile==2.0.0
+    # via openpyxl
+flexcache==0.3
+    # via pint
+flexparser==0.4
+    # via pint
+fsspec==2025.7.0
+    # via dask
+genno==1.28.2
+    # via ixmp
 gitdb==4.0.10
     # via gitpython
 gitpython==3.1.44
-    # via -r doc/requirements.in
+    # via -r requirements.in
 idna==3.7
     # via requests
 imagesize==1.4.1
     # via sphinx
+importlib-metadata==8.7.0
+    # via ixmp
+ixmp @ git+https://github.com/iiasa/ixmp.git@1730afd24238b4dad88a89c2d58d0fed2cd19dd0
+    # via -r requirements.in
 jinja2==3.1.6
     # via sphinx
+jpype1==1.6.0
+    # via ixmp
 latexcodec==2.0.1
     # via pybtex
+locket==1.0.0
+    # via partd
 markupsafe==2.1.3
     # via jinja2
 numpy==2.3.0
-    # via pandas
+    # via
+    #   dask
+    #   pandas
+    #   xarray
 numpydoc==1.8.0
-    # via -r doc/requirements.in
-packaging==23.1
-    # via sphinx
+    # via -r requirements.in
+openpyxl==3.1.5
+    # via ixmp
+packaging==25.0
+    # via
+    #   dask
+    #   jpype1
+    #   sphinx
+    #   xarray
 pandas==2.2.3
-    # via -r doc/requirements.in
+    # via
+    #   -r requirements.in
+    #   genno
+    #   ixmp
+    #   xarray
+partd==1.4.2
+    # via dask
+pint==0.24.4
+    # via
+    #   genno
+    #   ixmp
+platformdirs==4.3.8
+    # via
+    #   genno
+    #   pint
+pyarrow==21.0.0
+    # via pandas
 pybtex==0.24.0
     # via
     #   pybtex-docutils
@@ -49,7 +100,11 @@ python-dateutil==2.9.0.post0
 pytz==2025.2
     # via pandas
 pyyaml==6.0.1
-    # via pybtex
+    # via
+    #   dask
+    #   genno
+    #   ixmp
+    #   pybtex
 requests==2.32.4
     # via sphinx
 roman-numerals-py==3.1.0
@@ -67,18 +122,18 @@ snowballstemmer==2.2.0
     # via sphinx
 sphinx==8.2.3
     # via
-    #   -r doc/requirements.in
+    #   -r requirements.in
     #   numpydoc
     #   sphinx-rtd-theme
     #   sphinxcontrib-bibtex
     #   sphinxcontrib-jquery
     #   sphinxcontrib-serializinghtml
 sphinx-rtd-theme==3.0.2
-    # via -r doc/requirements.in
+    # via -r requirements.in
 sphinxcontrib-applehelp==2.0.0
     # via sphinx
 sphinxcontrib-bibtex==2.6.3
-    # via -r doc/requirements.in
+    # via -r requirements.in
 sphinxcontrib-devhelp==2.0.0
     # via sphinx
 sphinxcontrib-htmlhelp==2.1.0
@@ -93,7 +148,22 @@ sphinxcontrib-serializinghtml==1.1.9
     # via sphinx
 tabulate==0.9.0
     # via numpydoc
+toolz==1.0.0
+    # via
+    #   dask
+    #   partd
+typing-extensions==4.14.1
+    # via
+    #   flexcache
+    #   flexparser
+    #   pint
 tzdata==2025.2
     # via pandas
 urllib3==2.5.0
     # via requests
+xarray==2025.7.1
+    # via
+    #   genno
+    #   ixmp
+zipp==3.23.0
+    # via importlib-metadata

message_ix/cli.py

@@ -130,11 +130,3 @@ def dl(branch, tag, path):
 # Add subcommands
 main.add_command(message_ix.tools.add_year.cli.main)
 main.add_command(message_ix.tools.lp_diag.cli.main)
-
-try:
-    import message_ix.testing.nightly
-except ImportError:
-    # Dependencies of testing.nightly are missing; don't show the command
-    pass
-else:
-    main.add_command(message_ix.testing.nightly.cli)

message_ix/core.py

@@ -1,7 +1,7 @@
 import logging
 import os
 from collections.abc import Iterable, Mapping, Sequence
-from functools import lru_cache, partial
+from functools import lru_cache
 from itertools import chain, product, zip_longest
 from typing import Optional, TypeVar, Union
 from warnings import warn
@@ -12,8 +12,7 @@
 from ixmp.backend import ItemType
 from ixmp.backend.jdbc import JDBCBackend
 from ixmp.util import as_str_list, maybe_check_out, maybe_commit
-
-from message_ix.util.ixmp4 import on_ixmp4backend
+from ixmp.util.ixmp4 import is_ixmp4backend
 
 # from message_ix.util.scenario_data import PARAMETERS
 
@@ -210,7 +209,7 @@ def cat_list(self, name: str) -> list[str]:
         -------
         list of str
         """
-        return self._backend("cat_list", name)
+        return self.platform._backend.cat_list(self, name)
 
     def add_cat(self, name, cat, keys, is_unique=False):
         """Map elements from *keys* to category *cat* within set *name*.
@@ -227,7 +226,9 @@ def add_cat(self, name, cat, keys, is_unique=False):
             If `True`, then *cat* must have only one element. An exception is
             raised if *cat* already has an element, or if ``len(keys) > 1``.
         """
-        self._backend("cat_set_elements", name, str(cat), as_str_list(keys), is_unique)
+        self.platform._backend.cat_set_elements(
+            self, name, str(cat), as_str_list(keys), is_unique
+        )
 
     def cat(self, name, cat):
         """Return a list of all set elements mapped to a category.
@@ -247,7 +248,7 @@ def cat(self, name, cat):
         return list(
             map(
                 int if name == "year" else lambda v: v,
-                self._backend("cat_get_elements", name, cat),
+                self.platform._backend.cat_get_elements(self, name, cat),
             )
         )
 
@@ -257,16 +258,16 @@ def add_par(
         key_or_data: Optional[
             Union[int, str, Sequence[Union[int, str]], dict, pd.DataFrame]
         ] = None,
-        value=None,
-        unit: Optional[str] = None,
-        comment: Optional[str] = None,
+        value: Union[float, Iterable[float], None] = None,
+        unit: Union[str, Iterable[str], None] = None,
+        comment: Union[str, Iterable[str], None] = None,
     ) -> None:
         # ixmp.Scenario.add_par() is typed as accepting only str, but this method also
         # accepts int for "year"-like dimensions. Proxy the call to avoid type check
         # failures.
         # TODO Move this upstream, to ixmp
 
-        if on_ixmp4backend(self):
+        if is_ixmp4backend(self.platform._backend):
             from message_ix.util.scenario_setup import check_existence_of_units
 
             # Check for existence of required units
@@ -294,7 +295,13 @@ def add_par(
     def add_set(
         self,
         name: str,
-        key: Union[int, str, Sequence[Union[str, int]], dict, pd.DataFrame],
+        key: Union[
+            int,
+            str,
+            Iterable[object],
+            dict[str, Union[Sequence[int], Sequence[str]]],
+            pd.DataFrame,
+        ],
         comment: Union[str, Sequence[str], None] = None,
     ) -> None:
         # ixmp.Scenario.add_par() is typed as accepting only str, but this method also
@@ -841,10 +848,9 @@ def rename(self, name: str, mapping: Mapping[str, str], keep: bool = False) -> N
 
         # - Iterate over tuples of (item_name, ix_type); only those indexed by `name`.
         # - First all sets indexed sets; then all parameters.
-        items = partial(self.items, indexed_by=name, par_data=False)
         for item_name, ix_type in chain(
-            zip_longest(items(ItemType.SET), [], fillvalue="set"),
-            zip_longest(items(ItemType.PAR), [], fillvalue="par"),
+            zip_longest(self.items(ItemType.SET, indexed_by=name), [], fillvalue="set"),
+            zip_longest(self.items(ItemType.PAR, indexed_by=name), [], fillvalue="par"),
         ):
             # Identify some index names of this set; only those where the corresponding
             # index set is `name`