diff --git a/ixmp4/core/optimization/base.py b/ixmp4/core/optimization/base.py index 208435f7..68fe3b8d 100644 --- a/ixmp4/core/optimization/base.py +++ b/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) diff --git a/ixmp4/core/optimization/indexset.py b/ixmp4/core/optimization/indexset.py index e7fb8cc9..d73976c2 100644 --- a/ixmp4/core/optimization/indexset.py +++ b/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 - ) diff --git a/ixmp4/core/optimization/scalar.py b/ixmp4/core/optimization/scalar.py index fd5123c3..1586c2f4 100644 --- a/ixmp4/core/optimization/scalar.py +++ b/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"" -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 diff --git a/ixmp4/core/optimization/table.py b/ixmp4/core/optimization/table.py index 3afb5dc1..9ad75433 100644 --- a/ixmp4/core/optimization/table.py +++ b/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 diff --git a/ixmp4/data/abstract/optimization/base.py b/ixmp4/data/abstract/optimization/base.py index b10762db..b5018ece 100644 --- a/ixmp4/data/abstract/optimization/base.py +++ b/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: ... diff --git a/ixmp4/data/abstract/optimization/indexset.py b/ixmp4/data/abstract/optimization/indexset.py index cfacc96b..3a68c483 100644 --- a/ixmp4/data/abstract/optimization/indexset.py +++ b/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, diff --git a/ixmp4/data/abstract/optimization/scalar.py b/ixmp4/data/abstract/optimization/scalar.py index 0d0b52d8..7533c690 100644 --- a/ixmp4/data/abstract/optimization/scalar.py +++ b/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, diff --git a/ixmp4/data/abstract/optimization/table.py b/ixmp4/data/abstract/optimization/table.py index 7c2a71e3..d6721d2b 100644 --- a/ixmp4/data/abstract/optimization/table.py +++ b/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,