Skip to content

Commit 1043360

Browse files
committed
Implement dunder iadd, isub, imul methods for Quantity
imul is implemented only for `Percentage` because, all other multiplications that we have produce outputs of a new type, which is not possible with imul. With this usage like this becomes possible: ```python power = Power.from_watts(1000.0) power += Power.from_watts(3.2) ``` Signed-off-by: Sahas Subramanian <[email protected]>
1 parent 0c50e4b commit 1043360

File tree

2 files changed

+74
-0
lines changed

2 files changed

+74
-0
lines changed

src/frequenz/sdk/timeseries/_quantities.py

Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@
33

44
"""Types for holding quantities with units."""
55

6+
# pylint: disable=too-many-lines
7+
68
from __future__ import annotations
79

810
import math
@@ -237,6 +239,48 @@ def __mul__(self, percent: Percentage) -> Self:
237239
product._base_value = self._base_value * percent.as_fraction()
238240
return product
239241

242+
def __iadd__(self, other: Self) -> Self:
243+
"""Add another quantity to this one.
244+
245+
Args:
246+
other: The other quantity.
247+
248+
Returns:
249+
This quantity.
250+
"""
251+
if not type(other) is type(self):
252+
return NotImplemented
253+
self._base_value += other._base_value
254+
return self
255+
256+
def __isub__(self, other: Self) -> Self:
257+
"""Subtract another quantity from this one.
258+
259+
Args:
260+
other: The other quantity.
261+
262+
Returns:
263+
This quantity.
264+
"""
265+
if not type(other) is type(self):
266+
return NotImplemented
267+
self._base_value -= other._base_value
268+
return self
269+
270+
def __imul__(self, percent: Percentage) -> Self:
271+
"""Multiply this quantity by a percentage.
272+
273+
Args:
274+
percent: The percentage.
275+
276+
Returns:
277+
This quantity.
278+
"""
279+
if not isinstance(percent, Percentage):
280+
return NotImplemented
281+
self._base_value *= percent.as_fraction()
282+
return self
283+
240284
def __gt__(self, other: Self) -> bool:
241285
"""Return whether this quantity is greater than another.
242286

tests/timeseries/test_quantities.py

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,15 @@ def test_addition_subtraction() -> None:
105105
assert Fz1(1) - Fz2(1) # type: ignore
106106
assert excinfo.value.args[0] == "unsupported operand type(s) for -: 'Fz1' and 'Fz2'"
107107

108+
fz1 = Fz1(1.0)
109+
fz1 += Fz1(4.0)
110+
assert fz1 == Fz1(5.0)
111+
fz1 -= Fz1(9.0)
112+
assert fz1 == Fz1(-4.0)
113+
114+
with pytest.raises(TypeError) as excinfo:
115+
fz1 += Fz2(1.0) # type: ignore
116+
108117

109118
def test_comparison() -> None:
110119
"""Test the comparison of the quantities."""
@@ -394,6 +403,17 @@ def test_quantity_multiplied_with_precentage() -> None:
394403
assert energy * percentage == Energy.from_kilowatt_hours(6)
395404
assert percentage_ * percentage == Percentage.from_percent(25)
396405

406+
power *= percentage
407+
assert power == Power.from_watts(500.0)
408+
voltage *= percentage
409+
assert voltage == Voltage.from_volts(115.0)
410+
current *= percentage
411+
assert current == Current.from_amperes(1)
412+
energy *= percentage
413+
assert energy == Energy.from_kilowatt_hours(6)
414+
percentage_ *= percentage
415+
assert percentage_ == Percentage.from_percent(25)
416+
397417

398418
def test_invalid_multiplications() -> None:
399419
"""Test the multiplication of quantities with invalid quantities."""
@@ -405,19 +425,29 @@ def test_invalid_multiplications() -> None:
405425
for quantity in [power, voltage, current, energy]:
406426
with pytest.raises(TypeError):
407427
_ = power * quantity # type: ignore
428+
with pytest.raises(TypeError):
429+
power *= quantity # type: ignore
408430

409431
for quantity in [voltage, power, energy]:
410432
with pytest.raises(TypeError):
411433
_ = voltage * quantity # type: ignore
434+
with pytest.raises(TypeError):
435+
voltage *= quantity # type: ignore
412436

413437
for quantity in [current, power, energy]:
414438
with pytest.raises(TypeError):
415439
_ = current * quantity # type: ignore
440+
with pytest.raises(TypeError):
441+
current *= quantity # type: ignore
416442

417443
for quantity in [energy, power, voltage, current]:
418444
with pytest.raises(TypeError):
419445
_ = energy * quantity # type: ignore
446+
with pytest.raises(TypeError):
447+
energy *= quantity # type: ignore
420448

421449
for quantity in [power, voltage, current, energy, Percentage.from_percent(50)]:
422450
with pytest.raises(TypeError):
423451
_ = quantity * 200.0 # type: ignore
452+
with pytest.raises(TypeError):
453+
quantity *= 200.0 # type: ignore

0 commit comments

Comments
 (0)