DRY Parameter, Equation, Variable

* Spell out signature of create() to give proper type hints
iiasa · Dec 19, 2024 · abab5a5 · abab5a5
1 parent 632c385
commit abab5a5
Show file tree

Hide file tree

Showing 9 changed files with 59 additions and 92 deletions.
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
@@ -91,3 +91,6 @@ 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 create(self, name: str) -> IndexSet:
+        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/table.py b/ixmp4/core/optimization/table.py
@@ -98,3 +98,15 @@ 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
+
+    def create(
+        self,
+        name: str,
+        constrained_to_indexsets: list[str],
+        column_names: list[str] | None = None,
+    ) -> Table:
+        return super().create(
+            name=name,
+            constrained_to_indexsets=constrained_to_indexsets,
+            column_names=column_names,
+        )
diff --git a/ixmp4/core/optimization/variable.py b/ixmp4/core/optimization/variable.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 OptimizationVariable as VariableModel
 from ixmp4.data.abstract import Run
 from ixmp4.data.abstract.optimization import Column
 
+from .base import Lister, Retriever, Tabulator
+
 
 class Variable(BaseModelFacade):
     _model: VariableModel
@@ -110,12 +111,15 @@ def __str__(self) -> str:
         return f"<Variable {self.id} name={self.name}>"
 
 
-class VariableRepository(BaseFacade):
-    _run: Run
-
+class VariableRepository(
+    Retriever[Variable, VariableModel],
+    Lister[Variable, VariableModel],
+    Tabulator[Variable, VariableModel],
+):
     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.variables
+        self._model_type = Variable
 
     def create(
         self,
@@ -130,24 +134,3 @@ def create(
             column_names=column_names,
         )
         return Variable(_backend=self.backend, _model=model)
-
-    def get(self, name: str) -> Variable:
-        model = self.backend.optimization.variables.get(run_id=self._run.id, name=name)
-        return Variable(_backend=self.backend, _model=model)
-
-    def list(self, name: str | None = None) -> Iterable[Variable]:
-        variables = self.backend.optimization.variables.list(
-            run_id=self._run.id, name=name
-        )
-        return [
-            Variable(
-                _backend=self.backend,
-                _model=i,
-            )
-            for i in variables
-        ]
-
-    def tabulate(self, name: str | None = None) -> pd.DataFrame:
-        return self.backend.optimization.variables.tabulate(
-            run_id=self._run.id, name=name
-        )
diff --git a/ixmp4/data/abstract/optimization/base.py b/ixmp4/data/abstract/optimization/base.py
@@ -20,8 +20,7 @@
 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]
+    constrained_to_indexsets: str | list[str] | None
     column_names: list[str] | None
 
 

diff --git a/ixmp4/data/abstract/optimization/equation.py b/ixmp4/data/abstract/optimization/equation.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 EquationRepository(
+    BackendBaseRepository[Equation],
     base.Creator,
     base.Retriever,
     base.Enumerator,

diff --git a/ixmp4/data/abstract/optimization/parameter.py b/ixmp4/data/abstract/optimization/parameter.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 ParameterRepository(
+    BackendBaseRepository[Parameter],
     base.Creator,
     base.Retriever,
     base.Enumerator,

diff --git a/ixmp4/data/abstract/optimization/variable.py b/ixmp4/data/abstract/optimization/variable.py
@@ -12,6 +12,7 @@
 
 from .. import base
 from ..docs import DocsRepository
+from .base import BackendBaseRepository
 from .column import Column
 
 
@@ -38,6 +39,7 @@ def __str__(self) -> str:
 
 
 class VariableRepository(
+    BackendBaseRepository[Variable],
     base.Creator,
     base.Retriever,
     base.Enumerator,