DRY even more type hints using TypeAlias

iiasa · Nov 21, 2024 · e44a0c9 · e44a0c9
1 parent 1c45256
commit e44a0c9
Show file tree

Hide file tree

Showing 16 changed files with 96 additions and 318 deletions.
diff --git a/ixmp4/data/abstract/annotations.py b/ixmp4/data/abstract/annotations.py
@@ -14,6 +14,26 @@
 IntFilterAlias: TypeAlias = int | Iterable[int]
 StrFilterAlias: TypeAlias = str | Iterable[str]
 DefaultFilterAlias: TypeAlias = IntFilterAlias | StrFilterAlias
+OptimizationFilterAlias: TypeAlias = dict[str, DefaultFilterAlias | None]
+
+# NOTE If you want to be nitpicky, you could argue that timeseries have an additional
+# `variable` filter, which is not clear from this Alias used for both. However,
+# `variable` only adds more of the same types and we only use this for casting, so we
+# are fine *for now*.
+IamcFilterAlias: TypeAlias = dict[
+    str,
+    bool | DefaultFilterAlias | dict[str, DefaultFilterAlias] | None,
+]
+IamcObjectFilterAlias: TypeAlias = dict[
+    str,
+    DefaultFilterAlias
+    | dict[
+        str,
+        dict[str, DefaultFilterAlias | IamcFilterAlias],
+    ]
+    | bool
+    | None,
+]
 
 
 class HasIdFilter(TypedDict, total=False):

diff --git a/ixmp4/data/api/iamc/datapoint.py b/ixmp4/data/api/iamc/datapoint.py
@@ -1,5 +1,7 @@
 from datetime import datetime
-from typing import ClassVar, cast
+
+# TODO Use `type` instead of TypeAlias when dropping Python 3.11
+from typing import ClassVar, TypeAlias, cast
 
 import pandas as pd
 
@@ -26,6 +28,14 @@ class DataPoint(base.BaseModel):
     step_datetime: datetime | None
 
 
+JsonType: TypeAlias = dict[
+    str,
+    abstract.annotations.IntFilterAlias
+    | dict[str, bool | abstract.annotations.DefaultFilterAlias]
+    | None,
+]
+
+
 class DataPointRepository(
     base.Enumerator[DataPoint],
     base.BulkUpserter[DataPoint],
@@ -47,15 +57,7 @@ def list(
         **kwargs: Unpack[abstract.iamc.datapoint.EnumerateKwargs],
     ) -> list[DataPoint]:
         return super()._list(
-            json=cast(
-                dict[
-                    str,
-                    abstract.annotations.IntFilterAlias
-                    | dict[str, bool | abstract.annotations.DefaultFilterAlias]
-                    | None,
-                ],
-                kwargs,
-            ),
+            json=cast(JsonType, kwargs),
             params={
                 "join_parameters": join_parameters,
                 "join_runs": join_runs,
@@ -69,15 +71,7 @@ def tabulate(
         **kwargs: Unpack[abstract.iamc.datapoint.EnumerateKwargs],
     ) -> pd.DataFrame:
         return super()._tabulate(
-            json=cast(
-                dict[
-                    str,
-                    abstract.annotations.IntFilterAlias
-                    | dict[str, bool | abstract.annotations.DefaultFilterAlias]
-                    | None,
-                ],
-                kwargs,
-            ),
+            json=cast(JsonType, kwargs),
             params={
                 "join_parameters": join_parameters,
                 "join_runs": join_runs,

diff --git a/ixmp4/data/api/iamc/timeseries.py b/ixmp4/data/api/iamc/timeseries.py
@@ -45,39 +45,15 @@ def enumerate(
     def list(
         self, **kwargs: Unpack[abstract.iamc.timeseries.EnumerateKwargs]
     ) -> list[TimeSeries]:
-        json = cast(
-            dict[
-                str,
-                abstract.annotations.IntFilterAlias
-                | dict[
-                    str,
-                    bool
-                    | abstract.annotations.DefaultFilterAlias
-                    | dict[str, abstract.annotations.DefaultFilterAlias],
-                ],
-            ],
-            kwargs,
-        )
+        json = cast(abstract.annotations.IamcFilterAlias, kwargs)
         return super()._list(json=json)
 
     def tabulate(
         self,
         join_parameters: bool | None = None,
         **kwargs: Unpack[abstract.iamc.timeseries.EnumerateKwargs],
     ) -> pd.DataFrame:
-        json = cast(
-            dict[
-                str,
-                abstract.annotations.IntFilterAlias
-                | dict[
-                    str,
-                    bool
-                    | abstract.annotations.DefaultFilterAlias
-                    | dict[str, abstract.annotations.DefaultFilterAlias],
-                ],
-            ],
-            kwargs,
-        )
+        json = cast(abstract.annotations.IamcFilterAlias, kwargs)
         return super()._tabulate(json=json, params={"join_parameters": join_parameters})
 
     def get_by_id(self, id: int) -> TimeSeries:

diff --git a/ixmp4/data/api/iamc/variable.py b/ixmp4/data/api/iamc/variable.py
@@ -1,4 +1,3 @@
-from collections.abc import Mapping
 from datetime import datetime
 from typing import TYPE_CHECKING, ClassVar, cast
 
@@ -64,29 +63,11 @@ def enumerate(
     def list(
         self, **kwargs: Unpack[abstract.iamc.variable.EnumerateKwargs]
     ) -> list[Variable]:
-        json = cast(
-            dict[
-                str,
-                bool
-                | abstract.annotations.DefaultFilterAlias
-                | Mapping[str, abstract.annotations.DefaultFilterAlias]
-                | None,
-            ],
-            kwargs,
-        )
+        json = cast(abstract.annotations.IamcFilterAlias, kwargs)
         return super()._list(json=json)
 
     def tabulate(
         self, **kwargs: Unpack[abstract.iamc.variable.EnumerateKwargs]
     ) -> pd.DataFrame:
-        json = cast(
-            dict[
-                str,
-                bool
-                | abstract.annotations.DefaultFilterAlias
-                | Mapping[str, abstract.annotations.DefaultFilterAlias]
-                | None,
-            ],
-            kwargs,
-        )
+        json = cast(abstract.annotations.IamcFilterAlias, kwargs)
         return super()._tabulate(json=json)
diff --git a/ixmp4/data/api/meta.py b/ixmp4/data/api/meta.py
@@ -1,5 +1,7 @@
 from collections.abc import Iterable
-from typing import ClassVar, cast
+
+# TODO Use `type` instead of TypeAlias when dropping Python 3.11
+from typing import ClassVar, TypeAlias, cast
 
 import pandas as pd
 
@@ -23,6 +25,23 @@ class RunMetaEntry(base.BaseModel):
     value: abstract.StrictMetaValue
 
 
+# TODO This is tantalizingly close to the run JsonType, but not quite there.
+JsonType: TypeAlias = dict[
+    str,
+    bool
+    | float
+    | Iterable[float]
+    | abstract.annotations.DefaultFilterAlias
+    | dict[
+        str,
+        bool
+        | abstract.annotations.IntFilterAlias
+        | dict[str, abstract.annotations.DefaultFilterAlias],
+    ]
+    | None,
+]
+
+
 class RunMetaEntryRepository(
     base.Creator[RunMetaEntry],
     base.Retriever[RunMetaEntry],
@@ -60,47 +79,15 @@ def list(
         **kwargs: Unpack[abstract.meta.EnumerateKwargs],
     ) -> list[RunMetaEntry]:
         # base functions require dict, but TypedDict just inherits from Mapping
-        json = cast(
-            dict[
-                str,
-                bool
-                | float
-                | Iterable[float]
-                | abstract.annotations.DefaultFilterAlias
-                | dict[
-                    str,
-                    bool
-                    | abstract.annotations.IntFilterAlias
-                    | dict[str, abstract.annotations.DefaultFilterAlias],
-                ]
-                | None,
-            ],
-            kwargs,
-        )
+        json = cast(JsonType, kwargs)
         return super()._list(json=json, params={"join_run_index": join_run_index})
 
     def tabulate(
         self,
         join_run_index: bool | None = None,
         **kwargs: Unpack[abstract.meta.EnumerateKwargs],
     ) -> pd.DataFrame:
-        json = cast(
-            dict[
-                str,
-                bool
-                | float
-                | Iterable[float]
-                | abstract.annotations.DefaultFilterAlias
-                | dict[
-                    str,
-                    bool
-                    | abstract.annotations.IntFilterAlias
-                    | dict[str, abstract.annotations.DefaultFilterAlias],
-                ]
-                | None,
-            ],
-            kwargs,
-        )
+        json = cast(JsonType, kwargs)
         return super()._tabulate(json=json, params={"join_run_index": join_run_index})
 
     def bulk_upsert(self, df: pd.DataFrame) -> None:

diff --git a/ixmp4/data/api/model.py b/ixmp4/data/api/model.py
@@ -56,55 +56,13 @@ def enumerate(
         return super().enumerate(**kwargs)
 
     def list(self, **kwargs: Unpack[abstract.model.EnumerateKwargs]) -> list[Model]:
-        json = cast(
-            dict[
-                str,
-                abstract.annotations.DefaultFilterAlias
-                | dict[
-                    str,
-                    dict[
-                        str,
-                        abstract.annotations.DefaultFilterAlias
-                        | dict[
-                            str,
-                            bool
-                            | abstract.annotations.DefaultFilterAlias
-                            | dict[str, abstract.annotations.DefaultFilterAlias],
-                        ],
-                    ],
-                ]
-                | bool
-                | None,
-            ],
-            kwargs,
-        )
+        json = cast(abstract.annotations.IamcObjectFilterAlias, kwargs)
         return super()._list(json=json)
 
     def tabulate(
         self, **kwargs: Unpack[abstract.model.EnumerateKwargs]
     ) -> pd.DataFrame:
-        json = cast(
-            dict[
-                str,
-                abstract.annotations.DefaultFilterAlias
-                | dict[
-                    str,
-                    dict[
-                        str,
-                        abstract.annotations.DefaultFilterAlias
-                        | dict[
-                            str,
-                            bool
-                            | abstract.annotations.DefaultFilterAlias
-                            | dict[str, abstract.annotations.DefaultFilterAlias],
-                        ],
-                    ],
-                ]
-                | bool
-                | None,
-            ],
-            kwargs,
-        )
+        json = cast(abstract.annotations.IamcObjectFilterAlias, kwargs)
         return super()._tabulate(json=json)
 
 

diff --git a/ixmp4/data/api/optimization/equation.py b/ixmp4/data/api/optimization/equation.py
@@ -88,14 +88,14 @@ def list(
         self,
         **kwargs: Unpack[abstract.optimization.EnumerateKwargs],
     ) -> Iterable[Equation]:
-        json = cast(dict[str, abstract.annotations.DefaultFilterAlias | None], kwargs)
+        json = cast(abstract.annotations.OptimizationFilterAlias, kwargs)
         return super()._list(json=json)
 
     def tabulate(
         self,
         **kwargs: Unpack[abstract.optimization.EnumerateKwargs],
     ) -> pd.DataFrame:
-        json = cast(dict[str, abstract.annotations.DefaultFilterAlias | None], kwargs)
+        json = cast(abstract.annotations.OptimizationFilterAlias, kwargs)
         return super()._tabulate(json=json)
 
     def enumerate(

diff --git a/ixmp4/data/api/optimization/indexset.py b/ixmp4/data/api/optimization/indexset.py
@@ -73,14 +73,14 @@ def list(
         self,
         **kwargs: Unpack[abstract.optimization.EnumerateKwargs],
     ) -> list[IndexSet]:
-        json = cast(dict[str, abstract.annotations.DefaultFilterAlias | None], kwargs)
+        json = cast(abstract.annotations.OptimizationFilterAlias, kwargs)
         return super()._list(json=json)
 
     def tabulate(
         self,
         **kwargs: Unpack[abstract.optimization.EnumerateKwargs],
     ) -> pd.DataFrame:
-        json = cast(dict[str, abstract.annotations.DefaultFilterAlias | None], kwargs)
+        json = cast(abstract.annotations.OptimizationFilterAlias, kwargs)
         return super()._tabulate(json=json)
 
     def add_elements(

diff --git a/ixmp4/data/api/optimization/parameter.py b/ixmp4/data/api/optimization/parameter.py
@@ -85,14 +85,14 @@ def list(
         self,
         **kwargs: Unpack[abstract.optimization.EnumerateKwargs],
     ) -> Iterable[Parameter]:
-        json = cast(dict[str, abstract.annotations.DefaultFilterAlias | None], kwargs)
+        json = cast(abstract.annotations.OptimizationFilterAlias, kwargs)
         return super()._list(json=json)
 
     def tabulate(
         self,
         **kwargs: Unpack[abstract.optimization.EnumerateKwargs],
     ) -> pd.DataFrame:
-        json = cast(dict[str, abstract.annotations.DefaultFilterAlias | None], kwargs)
+        json = cast(abstract.annotations.OptimizationFilterAlias, kwargs)
         return super()._tabulate(json=json)
 
     def enumerate(

diff --git a/ixmp4/data/api/optimization/scalar.py b/ixmp4/data/api/optimization/scalar.py
@@ -86,13 +86,13 @@ def get_by_id(self, id: int) -> Scalar:
     def list(
         self, **kwargs: Unpack["abstract.optimization.scalar.EnumerateKwargs"]
     ) -> Iterable[Scalar]:
-        json = cast(dict[str, abstract.annotations.DefaultFilterAlias | None], kwargs)
+        json = cast(abstract.annotations.OptimizationFilterAlias, kwargs)
         return super()._list(json=json)
 
     def tabulate(
         self, **kwargs: Unpack["abstract.optimization.scalar.EnumerateKwargs"]
     ) -> pd.DataFrame:
-        json = cast(dict[str, abstract.annotations.DefaultFilterAlias | None], kwargs)
+        json = cast(abstract.annotations.OptimizationFilterAlias, kwargs)
         return super()._tabulate(json=json)
 
     def enumerate(

diff --git a/ixmp4/data/api/optimization/table.py b/ixmp4/data/api/optimization/table.py
@@ -84,13 +84,13 @@ def get_by_id(self, id: int) -> Table:
     def list(
         self, **kwargs: Unpack[abstract.optimization.EnumerateKwargs]
     ) -> Iterable[Table]:
-        json = cast(dict[str, abstract.annotations.DefaultFilterAlias | None], kwargs)
+        json = cast(abstract.annotations.OptimizationFilterAlias, kwargs)
         return super()._list(json=json)
 
     def tabulate(
         self, **kwargs: Unpack[abstract.optimization.EnumerateKwargs]
     ) -> pd.DataFrame:
-        json = cast(dict[str, abstract.annotations.DefaultFilterAlias | None], kwargs)
+        json = cast(abstract.annotations.OptimizationFilterAlias, kwargs)
         return super()._tabulate(json=json)
 
     def enumerate(