DRY optimization core layer (#96)

* Add abstraction layer for core optimization reusability * Add annotations
iiasa · Dec 19, 2024 · 0deac5b · 0deac5b
1 parent 26abf0e
commit 0deac5b
Show file tree

Hide file tree

Showing 16 changed files with 232 additions and 185 deletions.
diff --git a/ixmp4/core/optimization/base.py b/ixmp4/core/optimization/base.py
@@ -0,0 +1,94 @@
+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
+
+FacadeOptimizationModelType = TypeVar(
+    "FacadeOptimizationModelType", bound=BaseModelFacade
+)
+AbstractOptimizationModelType = TypeVar(
+    "AbstractOptimizationModelType", bound=abstract.BaseModel
+)
+
+
+class OptimizationBaseRepository(
+    BaseFacade, Generic[FacadeOptimizationModelType, AbstractOptimizationModelType]
+):
+    _run: abstract.Run
+    _backend_repository: abstract.BackendBaseRepository[AbstractOptimizationModelType]
+    _model_type: type[FacadeOptimizationModelType]
+
+    def __init__(self, _run: abstract.Run, **kwargs: Unpack["InitKwargs"]) -> None:
+        super().__init__(**kwargs)
+        self._run = _run
+
+
+class Creator(
+    OptimizationBaseRepository[
+        FacadeOptimizationModelType, AbstractOptimizationModelType
+    ],
+    abstract.Creator,
+):
+    def create(
+        self, name: str, **kwargs: Unpack["abstract.optimization.base.CreateKwargs"]
+    ) -> FacadeOptimizationModelType:
+        model = self._backend_repository.create(
+            run_id=self._run.id, name=name, **kwargs
+        )
+        return self._model_type(_backend=self.backend, _model=model)
+
+
+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[
+        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[
+        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[FacadeOptimizationModelType, AbstractOptimizationModelType],
+    Tabulator[FacadeOptimizationModelType, AbstractOptimizationModelType],
+    abstract.Enumerator,
+):
+    def enumerate(
+        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/equation.py b/ixmp4/core/optimization/equation.py
@@ -1,4 +1,3 @@
-from collections.abc import Iterable
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, ClassVar
 
@@ -10,12 +9,14 @@
 # TODO Import this from typing when dropping Python 3.11
 from typing_extensions import Unpack
 
-from ixmp4.core.base import BaseFacade, BaseModelFacade
+from ixmp4.core.base import BaseModelFacade
 from ixmp4.data.abstract import Docs as DocsModel
 from ixmp4.data.abstract import Equation as EquationModel
 from ixmp4.data.abstract import Run
 from ixmp4.data.abstract.optimization import Column
 
+from .base import Creator, Lister, Retriever, Tabulator
+
 
 class Equation(BaseModelFacade):
     _model: EquationModel
@@ -106,44 +107,25 @@ def __str__(self) -> str:
         return f"<Equation {self.id} name={self.name}>"
 
 
-class EquationRepository(BaseFacade):
-    _run: Run
-
+class EquationRepository(
+    Creator[Equation, EquationModel],
+    Retriever[Equation, EquationModel],
+    Lister[Equation, EquationModel],
+    Tabulator[Equation, EquationModel],
+):
     def __init__(self, _run: Run, **kwargs: Unpack["InitKwargs"]) -> None:
-        super().__init__(**kwargs)
-        self._run = _run
+        super().__init__(_run=_run, **kwargs)
+        self._backend_repository = self.backend.optimization.equations
+        self._model_type = Equation
 
     def create(
         self,
         name: str,
         constrained_to_indexsets: list[str],
         column_names: list[str] | None = None,
     ) -> Equation:
-        model = self.backend.optimization.equations.create(
+        return super().create(
             name=name,
-            run_id=self._run.id,
             constrained_to_indexsets=constrained_to_indexsets,
             column_names=column_names,
         )
-        return Equation(_backend=self.backend, _model=model)
-
-    def get(self, name: str) -> Equation:
-        model = self.backend.optimization.equations.get(run_id=self._run.id, name=name)
-        return Equation(_backend=self.backend, _model=model)
-
-    def list(self, name: str | None = None) -> Iterable[Equation]:
-        equations = self.backend.optimization.equations.list(
-            run_id=self._run.id, name=name
-        )
-        return [
-            Equation(
-                _backend=self.backend,
-                _model=i,
-            )
-            for i in equations
-        ]
-
-    def tabulate(self, name: str | None = None) -> pd.DataFrame:
-        return self.backend.optimization.equations.tabulate(
-            run_id=self._run.id, name=name
-        )
diff --git a/ixmp4/core/optimization/indexset.py b/ixmp4/core/optimization/indexset.py
@@ -4,16 +4,16 @@
 if TYPE_CHECKING:
     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 BaseFacade, BaseModelFacade
+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
+
 
 class IndexSet(BaseModelFacade):
     _model: IndexSetModel
@@ -81,39 +81,16 @@ def __str__(self) -> str:
         return f"<IndexSet {self.id} name={self.name}>"
 
 
-class IndexSetRepository(BaseFacade):
-    _run: Run
-
+class IndexSetRepository(
+    Creator[IndexSet, IndexSetModel],
+    Retriever[IndexSet, IndexSetModel],
+    Lister[IndexSet, IndexSetModel],
+    Tabulator[IndexSet, IndexSetModel],
+):
     def __init__(self, _run: Run, **kwargs: Unpack["InitKwargs"]) -> None:
-        super().__init__(**kwargs)
-        self._run = _run
+        super().__init__(_run=_run, **kwargs)
+        self._backend_repository = self.backend.optimization.indexsets
+        self._model_type = IndexSet
 
     def create(self, name: str) -> IndexSet:
-        indexset = self.backend.optimization.indexsets.create(
-            run_id=self._run.id,
-            name=name,
-        )
-        return IndexSet(_backend=self.backend, _model=indexset)
-
-    def get(self, name: str) -> IndexSet:
-        indexset = self.backend.optimization.indexsets.get(
-            run_id=self._run.id, name=name
-        )
-        return IndexSet(_backend=self.backend, _model=indexset)
-
-    def list(self, name: str | None = None) -> list[IndexSet]:
-        indexsets = self.backend.optimization.indexsets.list(
-            run_id=self._run.id, name=name
-        )
-        return [
-            IndexSet(
-                _backend=self.backend,
-                _model=i,
-            )
-            for i in indexsets
-        ]
-
-    def tabulate(self, name: str | None = None) -> pd.DataFrame:
-        return self.backend.optimization.indexsets.tabulate(
-            run_id=self._run.id, name=name
-        )
+        return super().create(name=name)
diff --git a/ixmp4/core/optimization/parameter.py b/ixmp4/core/optimization/parameter.py
@@ -1,4 +1,3 @@
-from collections.abc import Iterable
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, ClassVar
 
@@ -10,12 +9,14 @@
 # TODO Import this from typing when dropping Python 3.11
 from typing_extensions import Unpack
 
-from ixmp4.core.base import BaseFacade, BaseModelFacade
+from ixmp4.core.base import BaseModelFacade
 from ixmp4.data.abstract import Docs as DocsModel
 from ixmp4.data.abstract import Parameter as ParameterModel
 from ixmp4.data.abstract import Run, Unit
 from ixmp4.data.abstract.optimization import Column
 
+from .base import Creator, Lister, Retriever, Tabulator
+
 
 class Parameter(BaseModelFacade):
     _model: ParameterModel
@@ -99,44 +100,25 @@ def __str__(self) -> str:
         return f"<Parameter {self.id} name={self.name}>"
 
 
-class ParameterRepository(BaseFacade):
-    _run: Run
-
+class ParameterRepository(
+    Creator[Parameter, ParameterModel],
+    Retriever[Parameter, ParameterModel],
+    Lister[Parameter, ParameterModel],
+    Tabulator[Parameter, ParameterModel],
+):
     def __init__(self, _run: Run, **kwargs: Unpack["InitKwargs"]) -> None:
-        super().__init__(**kwargs)
-        self._run = _run
+        super().__init__(_run=_run, **kwargs)
+        self._backend_repository = self.backend.optimization.parameters
+        self._model_type = Parameter
 
     def create(
         self,
         name: str,
         constrained_to_indexsets: list[str],
         column_names: list[str] | None = None,
     ) -> Parameter:
-        model = self.backend.optimization.parameters.create(
+        return super().create(
             name=name,
-            run_id=self._run.id,
             constrained_to_indexsets=constrained_to_indexsets,
             column_names=column_names,
         )
-        return Parameter(_backend=self.backend, _model=model)
-
-    def get(self, name: str) -> Parameter:
-        model = self.backend.optimization.parameters.get(run_id=self._run.id, name=name)
-        return Parameter(_backend=self.backend, _model=model)
-
-    def list(self, name: str | None = None) -> Iterable[Parameter]:
-        parameters = self.backend.optimization.parameters.list(
-            run_id=self._run.id, name=name
-        )
-        return [
-            Parameter(
-                _backend=self.backend,
-                _model=i,
-            )
-            for i in parameters
-        ]
-
-    def tabulate(self, name: str | None = None) -> pd.DataFrame:
-        return self.backend.optimization.parameters.tabulate(
-            run_id=self._run.id, name=name
-        )
diff --git a/ixmp4/core/optimization/scalar.py b/ixmp4/core/optimization/scalar.py
@@ -1,22 +1,21 @@
-from collections.abc import Iterable
 from datetime import datetime
 from typing import TYPE_CHECKING, ClassVar
 
 if TYPE_CHECKING:
     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 BaseFacade, BaseModelFacade
+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
 
+from .base import Lister, Retriever, Tabulator
+
 
 class Scalar(BaseModelFacade):
     _model: ScalarModel
@@ -99,12 +98,15 @@ def __str__(self) -> str:
         return f"<Scalar {self.id} name={self.name}>"
 
 
-class ScalarRepository(BaseFacade):
-    _run: Run
-
+class ScalarRepository(
+    Retriever[Scalar, ScalarModel],
+    Lister[Scalar, ScalarModel],
+    Tabulator[Scalar, ScalarModel],
+):
     def __init__(self, _run: Run, **kwargs: Unpack["InitKwargs"]) -> None:
-        super().__init__(**kwargs)
-        self._run = _run
+        super().__init__(_run=_run, **kwargs)
+        self._backend_repository = self.backend.optimization.scalars
+        self._model_type = Scalar
 
     def create(self, name: str, value: float, unit: str | Unit | None = None) -> Scalar:
         if isinstance(unit, Unit):
@@ -127,22 +129,3 @@ def create(self, name: str, value: float, unit: str | Unit | None = None) -> Sca
                 "run.optimization.scalars.update()?"
             ) from e
         return Scalar(_backend=self.backend, _model=model)
-
-    def get(self, name: str) -> Scalar:
-        model = self.backend.optimization.scalars.get(run_id=self._run.id, name=name)
-        return Scalar(_backend=self.backend, _model=model)
-
-    def list(self, name: str | None = None) -> Iterable[Scalar]:
-        scalars = self.backend.optimization.scalars.list(run_id=self._run.id, name=name)
-        return [
-            Scalar(
-                _backend=self.backend,
-                _model=i,
-            )
-            for i in scalars
-        ]
-
-    def tabulate(self, name: str | None = None) -> pd.DataFrame:
-        return self.backend.optimization.scalars.tabulate(
-            run_id=self._run.id, name=name
-        )