Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 21 additions & 15 deletions examples/saved_scenarios_everywhere.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,9 @@
"id": "cba27484",
"metadata": {},
"source": [
"Testing notebook for interpreting everything as a saved scenario and delegating everywhere to scenario methods.\n"
"Testing notebook for loading mixed Scenario and Session types from Excel.\n",
"\n",
"This notebook demonstrates the simplified API where `Scenarios.from_excel()` loads both SavedScenario (from MyETM) and Session (from ETEngine) instances."
]
},
{
Expand Down Expand Up @@ -37,10 +39,9 @@
"outputs": [],
"source": [
"from pyetm.models.scenarios import Scenarios\n",
"from pyetm.models.sessions import Sessions\n",
"\n",
"scenarios = Scenarios.from_excel(\"../examples/inputs/example_input_excel.xlsx\")\n",
"sessions = Sessions.from_excel(\"../examples/inputs/example_input_excel.xlsx\")"
"# Load all scenarios (both SavedScenarios and Sessions) from the Excel file\n",
"mixed_scenarios = Scenarios.from_excel(\"../examples/inputs/example_input_excel.xlsx\")"
]
},
{
Expand All @@ -50,8 +51,9 @@
"metadata": {},
"outputs": [],
"source": [
"# Metadata\n",
"for scenario in scenarios:\n",
"# Iterate over all scenarios (both Scenario and Session instances)\n",
"# The delegation pattern means both types have the same interface\n",
"for scenario in mixed_scenarios:\n",
" print(f\"Title: {scenario.title}\")\n",
" print(f\"ID: {scenario.id}\")\n",
" print(f\"Area: {scenario.area_code}\")\n",
Expand All @@ -68,8 +70,9 @@
"metadata": {},
"outputs": [],
"source": [
"# Metadata\n",
"sessions.extend(scenarios.sessions)\n",
"# Access the underlying Session instances from all items\n",
"# For Scenario instances, this unwraps the session; for Session instances, returns them directly\n",
"sessions = mixed_scenarios.sessions\n",
"\n",
"for session in sessions:\n",
" print(f\"Title: {session.title}\")\n",
Expand All @@ -88,11 +91,13 @@
"metadata": {},
"outputs": [],
"source": [
"# Accessing the session data from the saved scenarios\n",
"for scenario in scenarios:\n",
" print(f\"Keep combatible: {scenario.keep_compatible}\")\n",
" print(f\"ID: {scenario.scenario_id}\")\n",
" print(f\"Inputs: {scenario.inputs.to_dataframe()}\")\n",
"# Accessing delegated properties works on all types\n",
"for scenario in mixed_scenarios:\n",
" print(f\"Keep compatible: {scenario.keep_compatible}\")\n",
" print(f\"ID: {scenario.id}\")\n",
" print(f\"Type: {'SavedScenario' if hasattr(scenario, 'scenario_id') and hasattr(scenario, 'session') else 'Session'}\")\n",
" # Uncomment to see inputs:\n",
" # print(f\"Inputs: {scenario.inputs.to_dataframe()}\")\n",
" print(\"\")"
]
},
Expand All @@ -111,8 +116,9 @@
"metadata": {},
"outputs": [],
"source": [
"# Export the scenarios to excel\n",
"# scenarios.to_excel(\"../examples/outputs/scenarios.xlsx\") # This will create scenarios.xlsx and scenarios_exports.xlsx (if you've set exports to true in the output config)."
"# Export all scenarios (both types) to Excel\n",
"# mixed_scenarios.to_excel(\"../examples/outputs/mixed_scenarios.xlsx\")\n",
"# This will create mixed_scenarios.xlsx (and mixed_scenarios_exports.xlsx if you've configured exports)"
]
}
],
Expand Down
69 changes: 45 additions & 24 deletions src/pyetm/models/scenarios.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
from __future__ import annotations
from os import PathLike
from pathlib import Path
from typing import Iterable, Iterator, List
from typing import Iterable, Iterator, List, Union
from pydantic import Field
from pyetm.models.session import Session
from pyetm.models.base import Base
Expand All @@ -10,37 +10,46 @@

class Scenarios(Base):
"""
A collection of SavedScenario objects.
A collection of Scenario and Session objects.

This class can hold both SavedScenario (Scenario) instances from MyETM
and Session instances from ETEngine. Since Scenario delegates all operations
to its underlying Session, both types can be used interchangeably.
"""

items: List[Scenario] = Field(default_factory=list)
items: List[Union[Scenario, Session]] = Field(default_factory=list)

def __iter__(self) -> Iterator[Scenario]:
def __iter__(self) -> Iterator[Union[Scenario, Session]]:
return iter(self.items)

def __len__(self) -> int:
return len(self.items)

def __getitem__(self, index: int) -> Scenario:
def __getitem__(self, index: int) -> Union[Scenario, Session]:
return self.items[index]

def add(self, *saved_scenarios: Scenario) -> None:
self.items.extend(saved_scenarios)
def add(self, *scenarios: Union[Scenario, Session]) -> None:
self.items.extend(scenarios)

def extend(self, saved_scenarios: Iterable[Scenario]) -> None:
self.items.extend(list(saved_scenarios))
def extend(self, scenarios: Iterable[Union[Scenario, Session]]) -> None:
self.items.extend(list(scenarios))

@property
def sessions(self) -> List["Session"]:
"""
Get the underlying ETEngine Scenario objects for all SavedScenarios.
Get the underlying ETEngine Session objects for all items.

For Scenario (SavedScenario) instances, returns the underlying session.
For Session instances, returns them directly.

Returns:
List of Scenario instances (the underlying sessions)
List of Session instances
"""
from pyetm.models.session import Session

return [saved.session for saved in self.items]
return [
item.session if isinstance(item, Scenario) else item for item in self.items
]

@classmethod
def load_many(cls, saved_scenario_ids: Iterable[int]) -> "Scenarios":
Expand All @@ -63,15 +72,21 @@ def load_many(cls, saved_scenario_ids: Iterable[int]) -> "Scenarios":

def to_excel(self, path: PathLike | str, **export_options) -> None:
"""
Export all saved scenarios to Excel.
Export all scenarios to Excel.

Exports both Scenario (SavedScenario) and Session instances.
For SavedScenario instances, the scenario_id column will contain
the MyETM SavedScenario ID. For Session instances, it will contain
the ETEngine session ID.

Note: This exports the underlying session data from each SavedScenario.
The scenario_id column will contain SavedScenario IDs (MyETM IDs).
Args:
path: Output path for the Excel file
**export_options: Additional export options to pass to ScenarioExcelService
"""
from pyetm.utils.scenario_excel_service import ScenarioExcelService

if not self.items:
raise ValueError("No saved scenarios to export")
raise ValueError("No scenarios to export")

resolved_path = Path(path).expanduser().resolve()
ScenarioExcelService.export_to_excel(
Expand All @@ -81,21 +96,27 @@ def to_excel(self, path: PathLike | str, **export_options) -> None:
@classmethod
def from_excel(cls, xlsx_path: PathLike | str) -> "Scenarios":
"""
Import SavedScenarios from Excel file.
Import all scenarios from Excel file.

Loads both SavedScenarios (Scenario instances from MyETM) and Sessions
(Session instances from ETEngine) based on the 'session' column value.
All scenarios are included regardless of type.

Only loads scenarios where the 'session' column is False or missing.
Scenarios with session=True are ignored.
Args:
xlsx_path: Path to Excel file

Returns:
Scenarios collection containing all loaded scenarios (both types)
"""
from pyetm.models.scenario_packer import ScenarioPacker

resolved_path = Path(xlsx_path).expanduser().resolve()

packer = ScenarioPacker.from_excel(str(resolved_path))
all_scenarios = list(packer._scenarios())
saved_scenarios = [s for s in all_scenarios if isinstance(s, Scenario)]

if not saved_scenarios:
print(f"No SavedScenarios found in Excel file: {resolved_path}")
if not all_scenarios:
print(f"No scenarios found in Excel file: {resolved_path}")

saved_scenarios.sort(key=lambda s: s.id if hasattr(s, "id") else 0)
return cls(items=saved_scenarios)
all_scenarios.sort(key=lambda s: s.id)
return cls(items=all_scenarios)
Loading