Add annotations

ixmp4/core/optimization/base.py

@@ -1,61 +1,94 @@
-from typing import Generic, TypeVar
+from typing import TYPE_CHECKING, Generic, TypeVar
+
+if TYPE_CHECKING:
+    from . import InitKwargs
+
 
 import pandas as pd
 
+# TODO Import this from typing when dropping Python 3.11
+from typing_extensions import TypedDict, Unpack
+
 from ixmp4.core.base import BaseFacade, BaseModelFacade
 from ixmp4.data import abstract
 
-OptimizationModelType = TypeVar("OptimizationModelType", bound=BaseModelFacade)
+FacadeOptimizationModelType = TypeVar(
+    "FacadeOptimizationModelType", bound=BaseModelFacade
+)
+AbstractOptimizationModelType = TypeVar(
+    "AbstractOptimizationModelType", bound=abstract.BaseModel
+)
 
 
-class OptimizationBaseRepository(BaseFacade, Generic[OptimizationModelType]):
+class OptimizationBaseRepository(
+    BaseFacade, Generic[FacadeOptimizationModelType, AbstractOptimizationModelType]
+):
     _run: abstract.Run
-    _backend_repository: abstract.BackendBaseRepository
-    _model_type: type[OptimizationModelType]
+    _backend_repository: abstract.BackendBaseRepository[AbstractOptimizationModelType]
+    _model_type: type[FacadeOptimizationModelType]
 
-    def __init__(self, _run: abstract.Run, *args, **kwargs) -> None:
-        super().__init__(*args, **kwargs)
+    def __init__(self, _run: abstract.Run, **kwargs: Unpack["InitKwargs"]) -> None:
+        super().__init__(**kwargs)
         self._run = _run
 
 
-class Creator(OptimizationBaseRepository[OptimizationModelType], abstract.Creator):
+class Creator(
+    OptimizationBaseRepository[
+        FacadeOptimizationModelType, AbstractOptimizationModelType
+    ],
+    abstract.Creator,
+):
     def create(
-        self,
-        name: str,
-        # TODO But how do we now show in core layer that e.g. Table needs these?
-        # constrained_to_indexsets: list[str],
-        # column_names: list[str] | None = None,
-        *args,
-        **kwargs,
-    ) -> OptimizationModelType:
+        self, name: str, **kwargs: Unpack["abstract.optimization.base.CreateKwargs"]
+    ) -> FacadeOptimizationModelType:
         model = self._backend_repository.create(
-            *args,
-            **dict(kwargs, name=name, run_id=self._run.id),
+            run_id=self._run.id, name=name, **kwargs
         )
         return self._model_type(_backend=self.backend, _model=model)
 
 
-class Retriever(OptimizationBaseRepository[OptimizationModelType], abstract.Retriever):
-    def get(self, name: str, *args, **kwargs) -> OptimizationModelType:
+class Retriever(
+    OptimizationBaseRepository[
+        FacadeOptimizationModelType, AbstractOptimizationModelType
+    ],
+    abstract.Retriever,
+):
+    def get(self, name: str) -> FacadeOptimizationModelType:
         model = self._backend_repository.get(run_id=self._run.id, name=name)
         return self._model_type(_backend=self.backend, _model=model)
 
 
-class Lister(OptimizationBaseRepository[OptimizationModelType], abstract.Lister):
-    def list(self, name: str | None = None) -> list[OptimizationModelType]:
+class Lister(
+    OptimizationBaseRepository[
+        FacadeOptimizationModelType, AbstractOptimizationModelType
+    ],
+    abstract.Lister,
+):
+    def list(self, name: str | None = None) -> list[FacadeOptimizationModelType]:
         models = self._backend_repository.list(run_id=self._run.id, name=name)
         return [self._model_type(_backend=self.backend, _model=m) for m in models]
 
 
-class Tabulator(OptimizationBaseRepository[OptimizationModelType], abstract.Tabulator):
+class Tabulator(
+    OptimizationBaseRepository[
+        FacadeOptimizationModelType, AbstractOptimizationModelType
+    ],
+    abstract.Tabulator,
+):
     def tabulate(self, name: str | None = None) -> pd.DataFrame:
         return self._backend_repository.tabulate(run_id=self._run.id, name=name)
 
 
+class EnumerateKwargs(TypedDict, total=False):
+    name: str | None
+
+
 class Enumerator(
-    Lister[OptimizationModelType], Tabulator[OptimizationModelType], abstract.Enumerator
+    Lister[FacadeOptimizationModelType, AbstractOptimizationModelType],
+    Tabulator[FacadeOptimizationModelType, AbstractOptimizationModelType],
+    abstract.Enumerator,
 ):
     def enumerate(
-        self, *args, table: bool = False, **kwargs
-    ) -> list[OptimizationModelType] | pd.DataFrame:
-        return self.tabulate(*args, **kwargs) if table else self.list(*args, **kwargs)
+        self, table: bool = False, **kwargs: Unpack[EnumerateKwargs]
+    ) -> list[FacadeOptimizationModelType] | pd.DataFrame:
+        return self.tabulate(**kwargs) if table else self.list(**kwargs)

ixmp4/core/optimization/indexset.py

@@ -2,16 +2,15 @@
 from typing import TYPE_CHECKING, ClassVar
 
 if TYPE_CHECKING:
-    pass
-
+    from . import InitKwargs
 
 # TODO Import this from typing when dropping Python 3.11
-
-import pandas as pd
+from typing_extensions import Unpack
 
 from ixmp4.core.base import BaseModelFacade
 from ixmp4.data.abstract import Docs as DocsModel
 from ixmp4.data.abstract import IndexSet as IndexSetModel
+from ixmp4.data.abstract import Run
 
 from .base import Creator, Lister, Retriever, Tabulator
 
@@ -83,16 +82,12 @@ def __str__(self) -> str:
 
 
 class IndexSetRepository(
-    Creator[IndexSet], Retriever[IndexSet], Lister[IndexSet], Tabulator[IndexSet]
+    Creator[IndexSet, IndexSetModel],
+    Retriever[IndexSet, IndexSetModel],
+    Lister[IndexSet, IndexSetModel],
+    Tabulator[IndexSet, IndexSetModel],
 ):
-    def __init__(self, *args, **kwargs) -> None:
-        super().__init__(*args, **kwargs)
+    def __init__(self, _run: Run, **kwargs: Unpack["InitKwargs"]) -> None:
+        super().__init__(_run=_run, **kwargs)
         self._backend_repository = self.backend.optimization.indexsets
         self._model_type = IndexSet
-
-    def tabulate(
-        self, name: str | None = None, include_data: bool = False
-    ) -> pd.DataFrame:
-        return self._backend_repository.tabulate(
-            run_id=self._run.id, name=name, include_data=include_data
-        )

ixmp4/core/optimization/scalar.py

@@ -2,14 +2,15 @@
 from typing import TYPE_CHECKING, ClassVar
 
 if TYPE_CHECKING:
-    pass
-
+    from . import InitKwargs
 
 # TODO Import this from typing when dropping Python 3.11
+from typing_extensions import Unpack
 
 from ixmp4.core.base import BaseModelFacade
 from ixmp4.core.unit import Unit
 from ixmp4.data.abstract import Docs as DocsModel
+from ixmp4.data.abstract import Run
 from ixmp4.data.abstract import Scalar as ScalarModel
 from ixmp4.data.abstract import Unit as UnitModel
 
@@ -97,9 +98,13 @@ def __str__(self) -> str:
         return f"<Scalar {self.id} name={self.name}>"
 
 
-class ScalarRepository(Retriever[Scalar], Lister[Scalar], Tabulator[Scalar]):
-    def __init__(self, *args, **kwargs) -> None:
-        super().__init__(*args, **kwargs)
+class ScalarRepository(
+    Retriever[Scalar, ScalarModel],
+    Lister[Scalar, ScalarModel],
+    Tabulator[Scalar, ScalarModel],
+):
+    def __init__(self, _run: Run, **kwargs: Unpack["InitKwargs"]) -> None:
+        super().__init__(_run=_run, **kwargs)
         self._backend_repository = self.backend.optimization.scalars
         self._model_type = Scalar
 

ixmp4/core/optimization/table.py

@@ -2,13 +2,16 @@
 from typing import TYPE_CHECKING, Any, ClassVar
 
 if TYPE_CHECKING:
-    pass
+    from . import InitKwargs
 
 import pandas as pd
 
 # TODO Import this from typing when dropping Python 3.11
+from typing_extensions import Unpack
+
 from ixmp4.core.base import BaseModelFacade
 from ixmp4.data.abstract import Docs as DocsModel
+from ixmp4.data.abstract import Run
 from ixmp4.data.abstract import Table as TableModel
 from ixmp4.data.abstract.optimization import Column
 
@@ -86,9 +89,12 @@ def __str__(self) -> str:
 
 
 class TableRepository(
-    Creator[Table], Retriever[Table], Lister[Table], Tabulator[Table]
+    Creator[Table, TableModel],
+    Retriever[Table, TableModel],
+    Lister[Table, TableModel],
+    Tabulator[Table, TableModel],
 ):
-    def __init__(self, *args, **kwargs) -> None:
-        super().__init__(*args, **kwargs)
+    def __init__(self, _run: Run, **kwargs: Unpack["InitKwargs"]) -> None:
+        super().__init__(_run=_run, **kwargs)
         self._backend_repository = self.backend.optimization.tables
         self._model_type = Table

ixmp4/data/abstract/optimization/base.py

@@ -1,26 +1,28 @@
-from typing import Generic, Iterable, Protocol, TypeVar
+from collections.abc import Iterable
+from typing import TYPE_CHECKING, Generic, Protocol, TypeVar
+
+if TYPE_CHECKING:
+    from . import EnumerateKwargs
 
 import pandas as pd
 
-# from ixmp4.data import types
-from .. import base
-from ..docs import DocsRepository
+# TODO Import this from typing when dropping Python 3.11
+from typing_extensions import TypedDict, Unpack
 
-# from .column import Column
+from ixmp4.data.abstract.unit import Unit
 
+from .. import base
+from ..docs import DocsRepository
 
-# TODO Currently not in use
-# class OptimizationBaseModel(base.BaseModel, Protocol):
-#     id: types.Integer
-#     name: types.String
-#     data: types.JsonDict
-#     columns: types.Mapped[list[Column]]
-#     run__id: types.Integer
-#     created_at: types.DateTime
-#     created_by: types.String
+BackendModelType = TypeVar("BackendModelType", bound=base.BaseModel, covariant=True)
 
 
-BackendModelType = TypeVar("BackendModelType", bound=base.BaseModel, covariant=True)
+class CreateKwargs(TypedDict, total=False):
+    value: float
+    unit: str | Unit | None
+    # TODO But how do we now show in core layer that e.g. Table needs these?
+    constrained_to_indexsets: list[str]
+    column_names: list[str] | None
 
 
 class BackendBaseRepository(
@@ -33,26 +35,13 @@ class BackendBaseRepository(
     docs: DocsRepository
 
     def create(
-        self,
-        run_id: int,
-        name: str,
-        # constrained_to_indexsets: list[str],
-        # column_names: list[str] | None = None,
-        *args,
-        **kwargs,
+        self, run_id: int, name: str, **kwargs: Unpack["CreateKwargs"]
     ) -> BackendModelType: ...
 
     def get(self, run_id: int, name: str) -> BackendModelType: ...
 
     def list(
-        self, *, name: str | None = None, **kwargs
+        self, **kwargs: Unpack["EnumerateKwargs"]
     ) -> Iterable[BackendModelType]: ...
 
-    def tabulate(self, *, name: str | None = None, **kwargs) -> pd.DataFrame: ...
-
-    # TODO Not needed as type hint in core layer:
-    # def get_by_id(self, id: int) -> BackendModelType: ...
-
-    # def add_data(
-    #     self, table_id: int, data: dict[str, Any] | pd.DataFrame
-    # ) -> None: ...
+    def tabulate(self, **kwargs: Unpack["EnumerateKwargs"]) -> pd.DataFrame: ...

ixmp4/data/abstract/optimization/indexset.py

@@ -12,6 +12,7 @@
 
 from .. import base
 from ..docs import DocsRepository
+from .base import BackendBaseRepository
 
 
 class IndexSet(base.BaseModel, Protocol):
@@ -36,6 +37,7 @@ def __str__(self) -> str:
 
 
 class IndexSetRepository(
+    BackendBaseRepository[IndexSet],
     base.Creator,
     base.Retriever,
     base.Enumerator,

ixmp4/data/abstract/optimization/scalar.py

@@ -19,6 +19,7 @@ class EnumerateKwargs(BaseEnumerateKwargs, HasUnitIdFilter, total=False): ...
 from .. import base
 from ..docs import DocsRepository
 from ..unit import Unit
+from .base import BackendBaseRepository
 
 
 class Scalar(base.BaseModel, Protocol):
@@ -45,6 +46,7 @@ def __str__(self) -> str:
 
 
 class ScalarRepository(
+    BackendBaseRepository[Scalar],
     base.Creator,
     base.Retriever,
     base.Enumerator,

ixmp4/data/abstract/optimization/table.py

@@ -13,6 +13,7 @@
 
 from .. import base
 from ..docs import DocsRepository
+from .base import BackendBaseRepository
 from .column import Column
 
 
@@ -39,6 +40,7 @@ def __str__(self) -> str:
 
 
 class TableRepository(
+    BackendBaseRepository[Table],
     base.Creator,
     base.Retriever,
     base.Enumerator,