Skip to content

Commit 633a59e

Browse files
FBumannclaude
andcommitted
refactor: make Comparison follow dict/Mapping protocol
Iteration now yields case names instead of (name, FlowSystem) pairs, matching Python's dict/Mapping convention. Added keys(), values(), and items() methods for explicit access. Users iterating pairs should switch to `for name, fs in comp.items():`. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent 5c55d36 commit 633a59e

2 files changed

Lines changed: 38 additions & 8 deletions

File tree

flixopt/comparison.py

Lines changed: 20 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
)
2020

2121
if TYPE_CHECKING:
22-
from collections.abc import Iterator
22+
from collections.abc import ItemsView, Iterator, KeysView, ValuesView
2323

2424
from .flow_system import FlowSystem
2525

@@ -224,14 +224,30 @@ def __getitem__(self, key: int | str) -> FlowSystem:
224224
return self._systems[idx]
225225
raise KeyError(f"Case '{key}' not found. Available: {self._names}")
226226

227-
def __iter__(self) -> Iterator[tuple[str, FlowSystem]]:
228-
"""Iterate over (name, FlowSystem) pairs."""
229-
yield from zip(self._names, self._systems, strict=True)
227+
def __iter__(self) -> Iterator[str]:
228+
"""Iterate over case names, matching the ``dict`` / ``Mapping`` protocol.
229+
230+
Use :meth:`items` for ``(name, FlowSystem)`` pairs or :meth:`values`
231+
for FlowSystems.
232+
"""
233+
return iter(self._names)
230234

231235
def __contains__(self, key: str) -> bool:
232236
"""Check if a case name exists."""
233237
return key in self._names
234238

239+
def keys(self) -> KeysView[str]:
240+
"""Return a view of case names, like :meth:`dict.keys`."""
241+
return self.flow_systems.keys()
242+
243+
def values(self) -> ValuesView[FlowSystem]:
244+
"""Return a view of FlowSystems, like :meth:`dict.values`."""
245+
return self.flow_systems.values()
246+
247+
def items(self) -> ItemsView[str, FlowSystem]:
248+
"""Return a view of ``(name, FlowSystem)`` pairs, like :meth:`dict.items`."""
249+
return self.flow_systems.items()
250+
235251
@property
236252
def flow_systems(self) -> dict[str, FlowSystem]:
237253
"""Access underlying FlowSystems as a dict mapping name → FlowSystem."""

tests/test_comparison.py

Lines changed: 18 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -212,11 +212,25 @@ def test_getitem_invalid_index_raises(self, optimized_base, optimized_with_chp):
212212
with pytest.raises(IndexError):
213213
_ = comp[99]
214214

215-
def test_iter(self, optimized_base, optimized_with_chp):
216-
"""Iteration yields (name, FlowSystem) pairs."""
215+
def test_iter_yields_names(self, optimized_base, optimized_with_chp):
216+
"""Iteration yields case names, matching the dict/Mapping protocol."""
217217
comp = fx.Comparison([optimized_base, optimized_with_chp])
218-
items = list(comp)
219-
assert items == [('Base', optimized_base), ('WithCHP', optimized_with_chp)]
218+
assert list(comp) == ['Base', 'WithCHP']
219+
220+
def test_keys(self, optimized_base, optimized_with_chp):
221+
"""keys() returns case names."""
222+
comp = fx.Comparison([optimized_base, optimized_with_chp])
223+
assert list(comp.keys()) == ['Base', 'WithCHP']
224+
225+
def test_values(self, optimized_base, optimized_with_chp):
226+
"""values() returns FlowSystems."""
227+
comp = fx.Comparison([optimized_base, optimized_with_chp])
228+
assert list(comp.values()) == [optimized_base, optimized_with_chp]
229+
230+
def test_items(self, optimized_base, optimized_with_chp):
231+
"""items() returns (name, FlowSystem) pairs without warning."""
232+
comp = fx.Comparison([optimized_base, optimized_with_chp])
233+
assert list(comp.items()) == [('Base', optimized_base), ('WithCHP', optimized_with_chp)]
220234

221235
def test_contains(self, optimized_base, optimized_with_chp):
222236
"""'in' operator checks for case name."""

0 commit comments

Comments
 (0)