Fix typing for composing constants with formula builders (#980)

Closes #951

Author: shsms

RELEASE_NOTES.md

@@ -17,4 +17,4 @@
 
 ## Bug Fixes
 
-<!-- Here goes notable bug fixes that are worth a special mention or explanation -->
+- Fixed a typing issue that occurs in some cases when composing formulas with constants.

src/frequenz/sdk/timeseries/formula_engine/_formula_engine.py

@@ -842,12 +842,12 @@ def _push(
 
     @overload
     def __add__(
-        self, other: _CompositionType1Phase
+        self, other: _CompositionType1Phase | QuantityT
     ) -> HigherOrderFormulaBuilder[QuantityT]: ...
 
     @overload
     def __add__(
-        self, other: _CompositionType3Phase | QuantityT
+        self, other: _CompositionType3Phase
     ) -> HigherOrderFormulaBuilder3Phase[QuantityT]: ...
 
     def __add__(
@@ -870,12 +870,12 @@ def __add__(
 
     @overload
     def __sub__(
-        self, other: _CompositionType1Phase
+        self, other: _CompositionType1Phase | QuantityT
     ) -> HigherOrderFormulaBuilder[QuantityT]: ...
 
     @overload
     def __sub__(
-        self, other: _CompositionType3Phase | QuantityT
+        self, other: _CompositionType3Phase
     ) -> HigherOrderFormulaBuilder3Phase[QuantityT]: ...
 
     def __sub__(
@@ -899,12 +899,12 @@ def __sub__(
 
     @overload
     def __mul__(
-        self, other: _CompositionType1Phase
+        self, other: _CompositionType1Phase | float
     ) -> HigherOrderFormulaBuilder[QuantityT]: ...
 
     @overload
     def __mul__(
-        self, other: _CompositionType3Phase | float
+        self, other: _CompositionType3Phase
     ) -> HigherOrderFormulaBuilder3Phase[QuantityT]: ...
 
     def __mul__(
@@ -928,12 +928,12 @@ def __mul__(
 
     @overload
     def __truediv__(
-        self, other: _CompositionType1Phase
+        self, other: _CompositionType1Phase | float
     ) -> HigherOrderFormulaBuilder[QuantityT]: ...
 
     @overload
     def __truediv__(
-        self, other: _CompositionType3Phase | float
+        self, other: _CompositionType3Phase
     ) -> HigherOrderFormulaBuilder3Phase[QuantityT]: ...
 
     def __truediv__(
@@ -957,12 +957,12 @@ def __truediv__(
 
     @overload
     def max(
-        self, other: _CompositionType1Phase
+        self, other: _CompositionType1Phase | QuantityT
     ) -> HigherOrderFormulaBuilder[QuantityT]: ...
 
     @overload
     def max(
-        self, other: _CompositionType3Phase | QuantityT
+        self, other: _CompositionType3Phase
     ) -> HigherOrderFormulaBuilder3Phase[QuantityT]: ...
 
     def max(
@@ -985,12 +985,12 @@ def max(
 
     @overload
     def min(
-        self, other: _CompositionType1Phase
+        self, other: _CompositionType1Phase | QuantityT
     ) -> HigherOrderFormulaBuilder[QuantityT]: ...
 
     @overload
     def min(
-        self, other: _CompositionType3Phase | QuantityT
+        self, other: _CompositionType3Phase
     ) -> HigherOrderFormulaBuilder3Phase[QuantityT]: ...
 
     def min(

tests/timeseries/_formula_engine/test_formula_composition.py

@@ -12,7 +12,7 @@
 from pytest_mock import MockerFixture
 
 from frequenz.sdk import microgrid
-from frequenz.sdk.timeseries._quantities import Power
+from frequenz.sdk.timeseries import Power, Sample
 
 from ..mock_microgrid import MockMicrogrid
 from .utils import get_resampled_stream
@@ -303,7 +303,10 @@ async def test_formula_composition_min_max_const(
             max_pow = await engine_max_rx.receive()
             assert max_pow and max_pow.value and max_pow.value.isclose(Power.zero())
 
-    async def test_formula_composition_constant(self, mocker: MockerFixture) -> None:
+    async def test_formula_composition_constant(  # pylint: disable=too-many-locals
+        self,
+        mocker: MockerFixture,
+    ) -> None:
         """Test the composition of formulas with constant values."""
         async with (
             MockMicrogrid(grid_meter=True, mocker=mocker) as mockgrid,
@@ -334,40 +337,64 @@ async def test_formula_composition_constant(self, mocker: MockerFixture) -> None
                 engine_div._stop  # pylint: disable=protected-access
             )
 
+            engine_composite = (
+                (
+                    (grid.power + Power.from_watts(50.0)) / 2.0
+                    + grid.power
+                    - Power.from_watts(20.0)
+                )
+                * 2.0
+            ).build("grid_power_composite")
+
             await mockgrid.mock_resampler.send_meter_power([100.0])
 
             # Test addition
-            grid_power_addition = await engine_add.new_receiver().receive()
+            grid_power_addition: Sample[Power] = (
+                await engine_add.new_receiver().receive()
+            )
             assert grid_power_addition.value is not None
             assert math.isclose(
                 grid_power_addition.value.as_watts(),
                 150.0,
             )
 
             # Test subtraction
-            grid_power_subtraction = await engine_sub.new_receiver().receive()
+            grid_power_subtraction: Sample[Power] = (
+                await engine_sub.new_receiver().receive()
+            )
             assert grid_power_subtraction.value is not None
             assert math.isclose(
                 grid_power_subtraction.value.as_watts(),
                 0.0,
             )
 
             # Test multiplication
-            grid_power_multiplication = await engine_mul.new_receiver().receive()
+            grid_power_multiplication: Sample[Power] = (
+                await engine_mul.new_receiver().receive()
+            )
             assert grid_power_multiplication.value is not None
             assert math.isclose(
                 grid_power_multiplication.value.as_watts(),
                 200.0,
             )
 
             # Test division
-            grid_power_division = await engine_div.new_receiver().receive()
+            grid_power_division: Sample[Power] = (
+                await engine_div.new_receiver().receive()
+            )
             assert grid_power_division.value is not None
             assert math.isclose(
                 grid_power_division.value.as_watts(),
                 50.0,
             )
 
+            # Test composite formula
+            grid_power_composite: Sample[Power] = (
+                await engine_composite.new_receiver().receive()
+            )
+            assert grid_power_composite.value is not None
+            assert math.isclose(grid_power_composite.value.as_watts(), 310.0)
+
             # Test multiplication with a Quantity
             with pytest.raises(RuntimeError):
                 engine_assert = (grid.power * Power.from_watts(2.0)).build(  # type: ignore