diff --git a/ixmp4/core/run.py b/ixmp4/core/run.py index dceddd56..57aa38c2 100644 --- a/ixmp4/core/run.py +++ b/ixmp4/core/run.py @@ -71,7 +71,7 @@ def meta(self, meta): self._meta._set(meta) def set_as_default(self): - """Sets this run as the default version for its `model` + `scenario` combination.""" + """Sets this run as default version for this `model - scenario` combination.""" self.backend.runs.set_as_default_version(self._model.id) def unset_as_default(self): @@ -80,20 +80,16 @@ def unset_as_default(self): class RunRepository(BaseFacade): - def list(self, default_only: bool = True) -> Iterable[Run]: + def list(self, default_only: bool = True, **kwargs) -> Iterable[Run]: return [ Run(_backend=self.backend, _model=r) - for r in self.backend.runs.list(default_only=default_only) + for r in self.backend.runs.list(default_only=default_only, **kwargs) ] - def tabulate(self, default_only: bool = True) -> pd.DataFrame: - runs = self.backend.runs.tabulate(default_only=default_only) - runs["model"] = runs["model__id"].map( - dict([(m.id, m.name) for m in self.backend.models.list()]) - ) - runs["scenario"] = runs["scenario__id"].map( - dict([(s.id, s.name) for s in self.backend.scenarios.list()]) - ) + def tabulate(self, default_only: bool = True, **kwargs) -> pd.DataFrame: + runs = self.backend.runs.tabulate(default_only=default_only, **kwargs) + runs["model"] = runs["model__id"].map(self.backend.models.map()) + runs["scenario"] = runs["scenario__id"].map(self.backend.scenarios.map()) return runs[["id", "model", "scenario", "version", "is_default"]] diff --git a/ixmp4/data/abstract/model.py b/ixmp4/data/abstract/model.py index 276a3eba..4e6b7d91 100644 --- a/ixmp4/data/abstract/model.py +++ b/ixmp4/data/abstract/model.py @@ -113,3 +113,13 @@ def tabulate( """ ... + + def map(self, *args, **kwargs) -> dict: + """Return a mapping of model-id to model-name. + + Returns + ------- + :class:`dict`: + A dictionary `id` -> `name` + """ + return dict([(m.id, m.name) for m in self.list(*args, **kwargs)]) diff --git a/ixmp4/data/abstract/scenario.py b/ixmp4/data/abstract/scenario.py index cce305cc..89a4b56b 100644 --- a/ixmp4/data/abstract/scenario.py +++ b/ixmp4/data/abstract/scenario.py @@ -105,3 +105,13 @@ def tabulate( - name """ ... + + def map(self, *args, **kwargs) -> dict: + """Return a mapping of scenario-id to scenario-name. + + Returns + ------- + :class:`dict` + A dictionary `id` -> `name` + """ + return dict([(s.id, s.name) for s in self.list(*args, **kwargs)]) diff --git a/tests/data/test_model.py b/tests/data/test_model.py index ef83527b..c9a9458a 100644 --- a/tests/data/test_model.py +++ b/tests/data/test_model.py @@ -56,6 +56,11 @@ def test_tabulate_model(self, test_mp): models.drop(columns=["created_at", "created_by"]), true_models ) + def test_map_model(self, test_mp): + test_mp.Run("Model 1", "Scenario", version="new") + test_mp.Run("Model 2", "Scenario", version="new") + assert test_mp.backend.models.map() == {1: "Model 1", 2: "Model 2"} + def test_filter_model(self, test_mp): run1, _ = create_filter_test_data(test_mp) diff --git a/tests/data/test_scenario.py b/tests/data/test_scenario.py index d6359d33..68a92b87 100644 --- a/tests/data/test_scenario.py +++ b/tests/data/test_scenario.py @@ -57,6 +57,12 @@ def test_tabulate_scenario(self, test_mp): scenarios.drop(columns=["created_at", "created_by"]), true_scenarios ) + def test_map_scenario(self, test_mp): + test_mp.Run("Model", "Scenario 1", version="new") + test_mp.Run("Model", "Scenario 2", version="new") + + assert test_mp.backend.scenarios.map() == {1: "Scenario 1", 2: "Scenario 2"} + def test_filter_scenario(self, test_mp): run1, _ = create_filter_test_data(test_mp)