Skip to content

Commit ca12ecb

Browse files
committed
Change min/max/coalesce to accept variadic args
Signed-off-by: Sahas Subramanian <[email protected]>
1 parent 8173981 commit ca12ecb

File tree

5 files changed

+62
-74
lines changed

5 files changed

+62
-74
lines changed

src/frequenz/sdk/timeseries/formulas/_formula.py

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -127,24 +127,24 @@ def __truediv__(self, other: float) -> FormulaBuilder[QuantityT]:
127127

128128
def coalesce(
129129
self,
130-
other: list[FormulaBuilder[QuantityT] | QuantityT | Formula[QuantityT]],
130+
*other: FormulaBuilder[QuantityT] | QuantityT | Formula[QuantityT],
131131
) -> FormulaBuilder[QuantityT]:
132132
"""Create a coalesce operation node."""
133-
return FormulaBuilder(self, self._create_method).coalesce(other)
133+
return FormulaBuilder(self, self._create_method).coalesce(*other)
134134

135135
def min(
136136
self,
137-
other: list[FormulaBuilder[QuantityT] | QuantityT | Formula[QuantityT]],
137+
*other: FormulaBuilder[QuantityT] | QuantityT | Formula[QuantityT],
138138
) -> FormulaBuilder[QuantityT]:
139139
"""Create a min operation node."""
140-
return FormulaBuilder(self, self._create_method).min(other)
140+
return FormulaBuilder(self, self._create_method).min(*other)
141141

142142
def max(
143143
self,
144-
other: list[FormulaBuilder[QuantityT] | QuantityT | Formula[QuantityT]],
144+
*other: FormulaBuilder[QuantityT] | QuantityT | Formula[QuantityT],
145145
) -> FormulaBuilder[QuantityT]:
146146
"""Create a max operation node."""
147-
return FormulaBuilder(self, self._create_method).max(other)
147+
return FormulaBuilder(self, self._create_method).max(*other)
148148

149149

150150
class FormulaBuilder(Generic[QuantityT]):
@@ -265,7 +265,7 @@ def __truediv__(
265265

266266
def coalesce(
267267
self,
268-
other: list[FormulaBuilder[QuantityT] | QuantityT | Formula[QuantityT]],
268+
*other: FormulaBuilder[QuantityT] | QuantityT | Formula[QuantityT],
269269
) -> FormulaBuilder[QuantityT]:
270270
"""Create a coalesce operation node."""
271271
right_nodes: list[AstNode[QuantityT]] = []
@@ -298,7 +298,7 @@ def coalesce(
298298

299299
def min(
300300
self,
301-
other: list[FormulaBuilder[QuantityT] | QuantityT | Formula[QuantityT]],
301+
*other: FormulaBuilder[QuantityT] | QuantityT | Formula[QuantityT],
302302
) -> FormulaBuilder[QuantityT]:
303303
"""Create a min operation node."""
304304
right_nodes: list[AstNode[QuantityT]] = []
@@ -331,7 +331,7 @@ def min(
331331

332332
def max(
333333
self,
334-
other: list[FormulaBuilder[QuantityT] | QuantityT | Formula[QuantityT]],
334+
*other: FormulaBuilder[QuantityT] | QuantityT | Formula[QuantityT],
335335
) -> FormulaBuilder[QuantityT]:
336336
"""Create a max operation node."""
337337
right_nodes: list[AstNode[QuantityT]] = []

src/frequenz/sdk/timeseries/formulas/_formula_3_phase.py

Lines changed: 33 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -114,38 +114,32 @@ def __truediv__(
114114

115115
def coalesce(
116116
self,
117-
other: list[
118-
Formula3PhaseBuilder[QuantityT]
119-
| Formula3Phase[QuantityT]
120-
| tuple[QuantityT, QuantityT, QuantityT]
121-
],
117+
*other: Formula3PhaseBuilder[QuantityT]
118+
| Formula3Phase[QuantityT]
119+
| tuple[QuantityT, QuantityT, QuantityT],
122120
) -> Formula3PhaseBuilder[QuantityT]:
123121
"""Coalesce the three-phase formula with a default value."""
124122
return Formula3PhaseBuilder(self, create_method=self._create_method).coalesce(
125-
other
123+
*other
126124
)
127125

128126
def min(
129127
self,
130-
other: list[
131-
Formula3PhaseBuilder[QuantityT]
132-
| Formula3Phase[QuantityT]
133-
| tuple[QuantityT, QuantityT, QuantityT]
134-
],
128+
*other: Formula3PhaseBuilder[QuantityT]
129+
| Formula3Phase[QuantityT]
130+
| tuple[QuantityT, QuantityT, QuantityT],
135131
) -> Formula3PhaseBuilder[QuantityT]:
136132
"""Get the minimum of the three-phase formula with other formulas."""
137-
return Formula3PhaseBuilder(self, create_method=self._create_method).min(other)
133+
return Formula3PhaseBuilder(self, create_method=self._create_method).min(*other)
138134

139135
def max(
140136
self,
141-
other: list[
142-
Formula3PhaseBuilder[QuantityT]
143-
| Formula3Phase[QuantityT]
144-
| tuple[QuantityT, QuantityT, QuantityT]
145-
],
137+
*other: Formula3PhaseBuilder[QuantityT]
138+
| Formula3Phase[QuantityT]
139+
| tuple[QuantityT, QuantityT, QuantityT],
146140
) -> Formula3PhaseBuilder[QuantityT]:
147141
"""Get the maximum of the three-phase formula with other formulas."""
148-
return Formula3PhaseBuilder(self, create_method=self._create_method).max(other)
142+
return Formula3PhaseBuilder(self, create_method=self._create_method).max(*other)
149143

150144

151145
class Formula3PhaseBuilder(Generic[QuantityT]):
@@ -284,16 +278,14 @@ def __truediv__(
284278

285279
def coalesce(
286280
self,
287-
others: list[
288-
Formula3PhaseBuilder[QuantityT]
289-
| Formula3Phase[QuantityT]
290-
| tuple[QuantityT, QuantityT, QuantityT]
291-
],
281+
*others: Formula3PhaseBuilder[QuantityT]
282+
| Formula3Phase[QuantityT]
283+
| tuple[QuantityT, QuantityT, QuantityT],
292284
) -> Formula3PhaseBuilder[QuantityT]:
293285
"""Coalesce the three-phase formula with a default value.
294286
295287
Args:
296-
others: The default value to use when the formula evaluates to None.
288+
*others: The other formulas or default values to coalesce with.
297289
298290
Returns:
299291
A new three-phase formula builder representing the coalesced formula.
@@ -341,26 +333,24 @@ def coalesce(
341333
right_nodes_phase_3.append(item.root[2])
342334
return Formula3PhaseBuilder(
343335
(
344-
self.root[0].coalesce(right_nodes_phase_1),
345-
self.root[1].coalesce(right_nodes_phase_2),
346-
self.root[2].coalesce(right_nodes_phase_3),
336+
self.root[0].coalesce(*right_nodes_phase_1),
337+
self.root[1].coalesce(*right_nodes_phase_2),
338+
self.root[2].coalesce(*right_nodes_phase_3),
347339
),
348340
create_method=self._create_method,
349341
sub_formulas=sub_formulas,
350342
)
351343

352344
def min(
353345
self,
354-
others: list[
355-
Formula3PhaseBuilder[QuantityT]
356-
| Formula3Phase[QuantityT]
357-
| tuple[QuantityT, QuantityT, QuantityT]
358-
],
346+
*others: Formula3PhaseBuilder[QuantityT]
347+
| Formula3Phase[QuantityT]
348+
| tuple[QuantityT, QuantityT, QuantityT],
359349
) -> Formula3PhaseBuilder[QuantityT]:
360350
"""Get the minimum of the three-phase formula with other formulas.
361351
362352
Args:
363-
others: The other formulas to compare with.
353+
*others: The other formulas or values to compare with.
364354
365355
Returns:
366356
A new three-phase formula builder representing the minimum.
@@ -408,26 +398,24 @@ def min(
408398
right_nodes_phase_3.append(item.root[2])
409399
return Formula3PhaseBuilder(
410400
(
411-
self.root[0].min(right_nodes_phase_1),
412-
self.root[1].min(right_nodes_phase_2),
413-
self.root[2].min(right_nodes_phase_3),
401+
self.root[0].min(*right_nodes_phase_1),
402+
self.root[1].min(*right_nodes_phase_2),
403+
self.root[2].min(*right_nodes_phase_3),
414404
),
415405
create_method=self._create_method,
416406
sub_formulas=sub_formulas,
417407
)
418408

419409
def max(
420410
self,
421-
others: list[
422-
Formula3PhaseBuilder[QuantityT]
423-
| Formula3Phase[QuantityT]
424-
| tuple[QuantityT, QuantityT, QuantityT]
425-
],
411+
*others: Formula3PhaseBuilder[QuantityT]
412+
| Formula3Phase[QuantityT]
413+
| tuple[QuantityT, QuantityT, QuantityT],
426414
) -> Formula3PhaseBuilder[QuantityT]:
427415
"""Get the maximum of the three-phase formula with other formulas.
428416
429417
Args:
430-
others: The other formulas to compare with.
418+
*others: The other formulas or values to compare with.
431419
432420
Returns:
433421
A new three-phase formula builder representing the maximum.
@@ -475,9 +463,9 @@ def max(
475463
right_nodes_phase_3.append(item.root[2])
476464
return Formula3PhaseBuilder(
477465
(
478-
self.root[0].max(right_nodes_phase_1),
479-
self.root[1].max(right_nodes_phase_2),
480-
self.root[2].max(right_nodes_phase_3),
466+
self.root[0].max(*right_nodes_phase_1),
467+
self.root[1].max(*right_nodes_phase_2),
468+
self.root[2].max(*right_nodes_phase_3),
481469
),
482470
create_method=self._create_method,
483471
sub_formulas=sub_formulas,

tests/timeseries/_formulas/test_formula_composition.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -218,13 +218,13 @@ async def test_formula_composition_min_max(self, mocker: MockerFixture) -> None:
218218
grid = microgrid.grid()
219219
stack.push_async_callback(grid.stop)
220220

221-
formula_min = grid.power.min([logical_meter.chp_power]).build(
221+
formula_min = grid.power.min(logical_meter.chp_power).build(
222222
"grid_power_min"
223223
)
224224
stack.push_async_callback(formula_min.stop)
225225
formula_min_rx = formula_min.new_receiver()
226226

227-
formula_max = grid.power.max([logical_meter.chp_power]).build(
227+
formula_max = grid.power.max(logical_meter.chp_power).build(
228228
"grid_power_max"
229229
)
230230
stack.push_async_callback(formula_max.stop)
@@ -305,11 +305,11 @@ async def test_formula_composition_min_max_const(
305305
grid = microgrid.grid()
306306
stack.push_async_callback(grid.stop)
307307

308-
formula_min = grid.power.min([Power.zero()]).build("grid_power_min")
308+
formula_min = grid.power.min(Power.zero()).build("grid_power_min")
309309
stack.push_async_callback(formula_min.stop)
310310
formula_min_rx = formula_min.new_receiver()
311311

312-
formula_max = grid.power.max([Power.zero()]).build("grid_power_max")
312+
formula_max = grid.power.max(Power.zero()).build("grid_power_max")
313313
stack.push_async_callback(formula_max.stop)
314314
formula_max_rx = formula_max.new_receiver()
315315

tests/timeseries/_formulas/test_formulas.py

Lines changed: 13 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -507,31 +507,31 @@ async def test_max(self) -> None:
507507
"""Test the max function."""
508508
await self.run_test(
509509
3,
510-
lambda c2, c4, c5: c2 + c4.max([c5]),
510+
lambda c2, c4, c5: c2 + c4.max(c5),
511511
"[l2]([0](#0) + MAX([1](#1), [2](#2)))",
512512
[
513513
([10.0, 12.0, 15.0], 25.0),
514514
],
515515
)
516516
await self.run_test(
517517
3,
518-
lambda c2, c4, c5: (c2 + c4).max([c5]),
518+
lambda c2, c4, c5: (c2 + c4).max(c5),
519519
"[l2](MAX([0](#0) + [1](#1), [2](#2)))",
520520
[
521521
([10.0, 12.0, 15.0], 22.0),
522522
],
523523
)
524524
await self.run_test(
525525
3,
526-
lambda c2, c4, c5: (c2 + c4).max([c5]),
526+
lambda c2, c4, c5: (c2 + c4).max(c5),
527527
"[l2](MAX([0](#0) + [1](#1), [2](#2)))",
528528
[
529529
([10.0, 12.0, 15.0], 22.0),
530530
],
531531
)
532532
await self.run_test(
533533
3,
534-
lambda c2, c4, c5: c2 + c4.max([c5]),
534+
lambda c2, c4, c5: c2 + c4.max(c5),
535535
"[l2]([0](#0) + MAX([1](#1), [2](#2)))",
536536
[
537537
([10.0, 12.0, 15.0], 25.0),
@@ -542,31 +542,31 @@ async def test_min(self) -> None:
542542
"""Test the min function."""
543543
await self.run_test(
544544
3,
545-
lambda c2, c4, c5: (c2 + c4).min([c5]),
545+
lambda c2, c4, c5: (c2 + c4).min(c5),
546546
"[l2](MIN([0](#0) + [1](#1), [2](#2)))",
547547
[
548548
([10.0, 12.0, 15.0], 15.0),
549549
],
550550
)
551551
await self.run_test(
552552
3,
553-
lambda c2, c4, c5: c2 + c4.min([c5]),
553+
lambda c2, c4, c5: c2 + c4.min(c5),
554554
"[l2]([0](#0) + MIN([1](#1), [2](#2)))",
555555
[
556556
([10.0, 12.0, 15.0], 22.0),
557557
],
558558
)
559559
await self.run_test(
560560
3,
561-
lambda c2, c4, c5: (c2 + c4).min([c5]),
561+
lambda c2, c4, c5: (c2 + c4).min(c5),
562562
"[l2](MIN([0](#0) + [1](#1), [2](#2)))",
563563
[
564564
([10.0, 2.0, 15.0], 12.0),
565565
],
566566
)
567567
await self.run_test(
568568
3,
569-
lambda c2, c4, c5: c2 + c4.min([c5]),
569+
lambda c2, c4, c5: c2 + c4.min(c5),
570570
"[l2]([0](#0) + MIN([1](#1), [2](#2)))",
571571
[
572572
([10.0, 12.0, 15.0], 22.0),
@@ -577,7 +577,7 @@ async def test_coalesce(self) -> None:
577577
"""Test the coalesce function."""
578578
await self.run_test(
579579
3,
580-
lambda c2, c4, c5: c2.coalesce([c4, c5]),
580+
lambda c2, c4, c5: c2.coalesce(c4, c5),
581581
"[l2](COALESCE([0](#0), [1](#1), [2](#2)))",
582582
[
583583
([None, 12.0, 15.0], None),
@@ -591,7 +591,7 @@ async def test_coalesce(self) -> None:
591591

592592
await self.run_test(
593593
3,
594-
lambda c2, c4, c5: (c2 * 5.0).coalesce([c4 / 2.0, c5]),
594+
lambda c2, c4, c5: (c2 * 5.0).coalesce(c4 / 2.0, c5),
595595
"[l2](COALESCE([0](#0) * 5.0, [1](#1) / 2.0, [2](#2)))",
596596
[
597597
([None, 12.0, 15.0], None),
@@ -610,7 +610,7 @@ async def test_min_max_coalesce(self) -> None:
610610
"""Test min and max functions in combination."""
611611
await self.run_test(
612612
3,
613-
lambda c2, c4, c5: c2.min([c4]).max([c5]),
613+
lambda c2, c4, c5: c2.min(c4).max(c5),
614614
"[l2](MAX(MIN([0](#0), [1](#1)), [2](#2)))",
615615
[
616616
([4.0, 6.0, 5.0], 5.0),
@@ -619,8 +619,8 @@ async def test_min_max_coalesce(self) -> None:
619619

620620
await self.run_test(
621621
3,
622-
lambda c2, c4, c5: c2.min([Quantity(0.0)])
623-
+ (c4 - c5).max([Quantity(0.0)]).coalesce([Quantity(0.0)]),
622+
lambda c2, c4, c5: c2.min(Quantity(0.0))
623+
+ (c4 - c5).max(Quantity(0.0)).coalesce(Quantity(0.0)),
624624
"[l2](MIN([0](#0), 0.0) + COALESCE(MAX([1](#1) - [2](#2), 0.0), 0.0))",
625625
[
626626
([4.0, 6.0, 5.0], 1.0),

tests/timeseries/_formulas/test_formulas_3_phase.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -140,7 +140,7 @@ async def test_composition(self) -> None:
140140
([(-5.0, 10.0, 15.0), (10.0, 5.0, -20.0)], (10.0, 10.0, 15.0)),
141141
([(None, 2.0, 3.0), (2.0, None, 4.0)], (None, None, 4.0)),
142142
],
143-
lambda f1, f2: f1.max([f2]),
143+
lambda f1, f2: f1.max(f2),
144144
)
145145

146146
await self.run_test(
@@ -149,7 +149,7 @@ async def test_composition(self) -> None:
149149
([(-5.0, 10.0, 15.0), (10.0, 5.0, -20.0)], (-5.0, 5.0, -20.0)),
150150
([(None, 2.0, 3.0), (2.0, None, 4.0)], (None, None, 3.0)),
151151
],
152-
lambda f1, f2: f1.min([f2]),
152+
lambda f1, f2: f1.min(f2),
153153
)
154154

155155
await self.run_test(
@@ -161,6 +161,6 @@ async def test_composition(self) -> None:
161161
([(None, 2.0, 3.0), (2.0, None, 4.0)], (2.0, 2.0, 3.0)),
162162
],
163163
lambda f1, f2: f1.coalesce(
164-
[f2, (Quantity.zero(), Quantity.zero(), Quantity.zero())]
164+
f2, (Quantity.zero(), Quantity.zero(), Quantity.zero())
165165
),
166166
)

0 commit comments

Comments
 (0)