Skip to content

Commit d6580af

Browse files
shsmsshsms
committed
Update FormulaEnginePool to have multiple strongly-typed functions
The previous design used a single dictionary `_engines` to store engines of different types. This is fixed now. Signed-off-by: Sahas Subramanian <[email protected]>
1 parent e00cc91 commit d6580af

File tree

4 files changed

+80
-35
lines changed

4 files changed

+80
-35
lines changed

‎src/frequenz/sdk/timeseries/_formula_engine/_formula_engine_pool.py‎

Lines changed: 59 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -5,13 +5,13 @@
55

66
from __future__ import annotations
77

8-
from typing import TYPE_CHECKING, Any, Type
8+
from typing import TYPE_CHECKING, Type
99

1010
from frequenz.channels import Sender
1111

1212
from ...actor import ChannelRegistry, ComponentMetricRequest
1313
from ...microgrid.component import ComponentMetricId
14-
from .._quantities import Quantity
14+
from .._quantities import Current, Power, Quantity
1515
from ._formula_generators._formula_generator import (
1616
FormulaGenerator,
1717
FormulaGeneratorConfig,
@@ -47,7 +47,9 @@ def __init__(
4747
self._namespace = namespace
4848
self._channel_registry = channel_registry
4949
self._resampler_subscription_sender = resampler_subscription_sender
50-
self._engines: dict[str, "FormulaEngine[Any]|FormulaEngine3Phase[Any]"] = {}
50+
self._string_engines: dict[str, FormulaEngine[Quantity]] = {}
51+
self._power_engines: dict[str, FormulaEngine[Power]] = {}
52+
self._current_engines: dict[str, FormulaEngine3Phase[Current]] = {}
5153

5254
def from_string(
5355
self,
@@ -68,8 +70,8 @@ def from_string(
6870
A FormulaReceiver that streams values with the formulas applied.
6971
"""
7072
channel_key = formula + component_metric_id.value
71-
if channel_key in self._engines:
72-
return self._engines[channel_key] # type: ignore
73+
if channel_key in self._string_engines:
74+
return self._string_engines[channel_key]
7375

7476
builder = ResampledFormulaBuilder(
7577
self._namespace,
@@ -80,16 +82,16 @@ def from_string(
8082
Quantity,
8183
)
8284
formula_engine = builder.from_string(formula, nones_are_zeros)
83-
self._engines[channel_key] = formula_engine
85+
self._string_engines[channel_key] = formula_engine
8486

8587
return formula_engine
8688

87-
def from_generator(
89+
def from_power_formula_generator(
8890
self,
8991
channel_key: str,
90-
generator: "Type[FormulaGenerator[Any]]",
92+
generator: Type[FormulaGenerator[Power]],
9193
config: FormulaGeneratorConfig = FormulaGeneratorConfig(),
92-
) -> FormulaEngine[Any] | FormulaEngine3Phase[Any]:
94+
) -> FormulaEngine[Power]:
9395
"""Get a receiver for a formula from a generator.
9496
9597
Args:
@@ -101,19 +103,62 @@ def from_generator(
101103
A FormulaReceiver or a FormulaReceiver3Phase instance based on what the
102104
FormulaGenerator returns.
103105
"""
104-
if channel_key in self._engines:
105-
return self._engines[channel_key]
106+
from ._formula_engine import ( # pylint: disable=import-outside-toplevel
107+
FormulaEngine,
108+
)
109+
110+
if channel_key in self._power_engines:
111+
return self._power_engines[channel_key]
112+
113+
engine = generator(
114+
self._namespace,
115+
self._channel_registry,
116+
self._resampler_subscription_sender,
117+
config,
118+
).generate()
119+
assert isinstance(engine, FormulaEngine)
120+
self._power_engines[channel_key] = engine
121+
return engine
122+
123+
def from_3_phase_current_formula_generator(
124+
self,
125+
channel_key: str,
126+
generator: "Type[FormulaGenerator[Current]]",
127+
config: FormulaGeneratorConfig = FormulaGeneratorConfig(),
128+
) -> FormulaEngine3Phase[Current]:
129+
"""Get a receiver for a formula from a generator.
130+
131+
Args:
132+
channel_key: A string to uniquely identify the formula.
133+
generator: A formula generator.
134+
config: config to initialize the formula generator with.
135+
136+
Returns:
137+
A FormulaReceiver or a FormulaReceiver3Phase instance based on what the
138+
FormulaGenerator returns.
139+
"""
140+
from ._formula_engine import ( # pylint: disable=import-outside-toplevel
141+
FormulaEngine3Phase,
142+
)
143+
144+
if channel_key in self._current_engines:
145+
return self._current_engines[channel_key]
106146

107147
engine = generator(
108148
self._namespace,
109149
self._channel_registry,
110150
self._resampler_subscription_sender,
111151
config,
112152
).generate()
113-
self._engines[channel_key] = engine
153+
assert isinstance(engine, FormulaEngine3Phase)
154+
self._current_engines[channel_key] = engine
114155
return engine
115156

116157
async def stop(self) -> None:
117158
"""Stop all formula engines in the pool."""
118-
for engine in self._engines.values():
119-
await engine._stop() # pylint: disable=protected-access
159+
for string_engine in self._string_engines.values():
160+
await string_engine._stop() # pylint: disable=protected-access
161+
for power_engine in self._power_engines.values():
162+
await power_engine._stop() # pylint: disable=protected-access
163+
for current_engine in self._current_engines.values():
164+
await current_engine._stop() # pylint: disable=protected-access

‎src/frequenz/sdk/timeseries/battery_pool/battery_pool.py‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -253,7 +253,7 @@ def power(self) -> FormulaEngine[Power]:
253253
A FormulaEngine that will calculate and stream the total power of all
254254
batteries in the pool.
255255
"""
256-
engine = self._formula_pool.from_generator(
256+
engine = self._formula_pool.from_power_formula_generator(
257257
"battery_pool_power",
258258
BatteryPowerFormula,
259259
FormulaGeneratorConfig(
@@ -280,7 +280,7 @@ def production_power(self) -> FormulaEngine[Power]:
280280
A FormulaEngine that will calculate and stream the total production power of
281281
all batteries in the pool.
282282
"""
283-
engine = self._formula_pool.from_generator(
283+
engine = self._formula_pool.from_power_formula_generator(
284284
"battery_pool_production_power",
285285
BatteryPowerFormula,
286286
FormulaGeneratorConfig(
@@ -304,10 +304,10 @@ def consumption_power(self) -> FormulaEngine[Power]:
304304
method.
305305
306306
Returns:
307-
A FormulaEngine that will calculate and stream the total consumption power of
308-
all batteries in the pool.
307+
A FormulaEngine that will calculate and stream the total consumption
308+
power of all batteries in the pool.
309309
"""
310-
engine = self._formula_pool.from_generator(
310+
engine = self._formula_pool.from_power_formula_generator(
311311
"battery_pool_consumption_power",
312312
BatteryPowerFormula,
313313
FormulaGeneratorConfig(

‎src/frequenz/sdk/timeseries/ev_charger_pool/_ev_charger_pool.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -123,7 +123,7 @@ def current(self) -> FormulaEngine3Phase[Current]:
123123
A FormulaEngine that will calculate and stream the total current of all EV
124124
Chargers.
125125
"""
126-
engine = self._formula_pool.from_generator(
126+
engine = self._formula_pool.from_3_phase_current_formula_generator(
127127
"ev_charger_total_current",
128128
EVChargerCurrentFormula,
129129
FormulaGeneratorConfig(component_ids=self._component_ids),
@@ -147,7 +147,7 @@ def power(self) -> FormulaEngine[Power]:
147147
A FormulaEngine that will calculate and stream the total power of all EV
148148
Chargers.
149149
"""
150-
engine = self._formula_pool.from_generator(
150+
engine = self._formula_pool.from_power_formula_generator(
151151
"ev_charger_power",
152152
EVChargerPowerFormula,
153153
FormulaGeneratorConfig(
@@ -174,7 +174,7 @@ def production_power(self) -> FormulaEngine[Power]:
174174
A FormulaEngine that will calculate and stream the production power of all
175175
EV Chargers.
176176
"""
177-
engine = self._formula_pool.from_generator(
177+
engine = self._formula_pool.from_power_formula_generator(
178178
"ev_charger_production_power",
179179
EVChargerPowerFormula,
180180
FormulaGeneratorConfig(
@@ -201,7 +201,7 @@ def consumption_power(self) -> FormulaEngine[Power]:
201201
A FormulaEngine that will calculate and stream the consumption power of all
202202
EV Chargers.
203203
"""
204-
engine = self._formula_pool.from_generator(
204+
engine = self._formula_pool.from_power_formula_generator(
205205
"ev_charger_consumption_power",
206206
EVChargerPowerFormula,
207207
FormulaGeneratorConfig(

‎src/frequenz/sdk/timeseries/logical_meter/_logical_meter.py‎

Lines changed: 12 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -163,7 +163,7 @@ def grid_power(self) -> FormulaEngine[Power]:
163163
Returns:
164164
A FormulaEngine that will calculate and stream grid power.
165165
"""
166-
engine = self._formula_pool.from_generator(
166+
engine = self._formula_pool.from_power_formula_generator(
167167
"grid_power",
168168
GridPowerFormula,
169169
)
@@ -185,7 +185,7 @@ def grid_consumption_power(self) -> FormulaEngine[Power]:
185185
Returns:
186186
A FormulaEngine that will calculate and stream grid consumption power.
187187
"""
188-
engine = self._formula_pool.from_generator(
188+
engine = self._formula_pool.from_power_formula_generator(
189189
"grid_consumption_power",
190190
GridPowerFormula,
191191
FormulaGeneratorConfig(formula_type=FormulaType.CONSUMPTION),
@@ -208,7 +208,7 @@ def grid_production_power(self) -> FormulaEngine[Power]:
208208
Returns:
209209
A FormulaEngine that will calculate and stream grid production power.
210210
"""
211-
engine = self._formula_pool.from_generator(
211+
engine = self._formula_pool.from_power_formula_generator(
212212
"grid_production_power",
213213
GridPowerFormula,
214214
FormulaGeneratorConfig(formula_type=FormulaType.PRODUCTION),
@@ -231,7 +231,7 @@ def grid_current(self) -> FormulaEngine3Phase[Current]:
231231
Returns:
232232
A FormulaEngine that will calculate and stream grid current.
233233
"""
234-
engine = self._formula_pool.from_generator(
234+
engine = self._formula_pool.from_3_phase_current_formula_generator(
235235
"grid_current",
236236
GridCurrentFormula,
237237
)
@@ -256,7 +256,7 @@ def consumer_power(self) -> FormulaEngine[Power]:
256256
Returns:
257257
A FormulaEngine that will calculate and stream consumer power.
258258
"""
259-
engine = self._formula_pool.from_generator(
259+
engine = self._formula_pool.from_power_formula_generator(
260260
"consumer_power",
261261
ConsumerPowerFormula,
262262
)
@@ -281,7 +281,7 @@ def producer_power(self) -> FormulaEngine[Power]:
281281
Returns:
282282
A FormulaEngine that will calculate and stream producer power.
283283
"""
284-
engine = self._formula_pool.from_generator(
284+
engine = self._formula_pool.from_power_formula_generator(
285285
"producer_power",
286286
ProducerPowerFormula,
287287
)
@@ -303,7 +303,7 @@ def pv_power(self) -> FormulaEngine[Power]:
303303
Returns:
304304
A FormulaEngine that will calculate and stream PV total power.
305305
"""
306-
engine = self._formula_pool.from_generator(
306+
engine = self._formula_pool.from_power_formula_generator(
307307
"pv_power",
308308
PVPowerFormula,
309309
FormulaGeneratorConfig(formula_type=FormulaType.PASSIVE_SIGN_CONVENTION),
@@ -326,7 +326,7 @@ def pv_production_power(self) -> FormulaEngine[Power]:
326326
Returns:
327327
A FormulaEngine that will calculate and stream PV power production.
328328
"""
329-
engine = self._formula_pool.from_generator(
329+
engine = self._formula_pool.from_power_formula_generator(
330330
"pv_production_power",
331331
PVPowerFormula,
332332
FormulaGeneratorConfig(formula_type=FormulaType.PRODUCTION),
@@ -349,7 +349,7 @@ def pv_consumption_power(self) -> FormulaEngine[Power]:
349349
Returns:
350350
A FormulaEngine that will calculate and stream PV power consumption.
351351
"""
352-
engine = self._formula_pool.from_generator(
352+
engine = self._formula_pool.from_power_formula_generator(
353353
"pv_consumption_power",
354354
PVPowerFormula,
355355
FormulaGeneratorConfig(formula_type=FormulaType.CONSUMPTION),
@@ -372,7 +372,7 @@ def chp_power(self) -> FormulaEngine[Power]:
372372
Returns:
373373
A FormulaEngine that will calculate and stream CHP power production.
374374
"""
375-
engine = self._formula_pool.from_generator(
375+
engine = self._formula_pool.from_power_formula_generator(
376376
"chp_power",
377377
CHPPowerFormula,
378378
FormulaGeneratorConfig(formula_type=FormulaType.PASSIVE_SIGN_CONVENTION),
@@ -395,7 +395,7 @@ def chp_production_power(self) -> FormulaEngine[Power]:
395395
Returns:
396396
A FormulaEngine that will calculate and stream CHP power production.
397397
"""
398-
engine = self._formula_pool.from_generator(
398+
engine = self._formula_pool.from_power_formula_generator(
399399
"chp_production_power",
400400
CHPPowerFormula,
401401
FormulaGeneratorConfig(
@@ -420,7 +420,7 @@ def chp_consumption_power(self) -> FormulaEngine[Power]:
420420
Returns:
421421
A FormulaEngine that will calculate and stream CHP power consumption.
422422
"""
423-
engine = self._formula_pool.from_generator(
423+
engine = self._formula_pool.from_power_formula_generator(
424424
"chp_consumption_power",
425425
CHPPowerFormula,
426426
FormulaGeneratorConfig(

0 commit comments

Comments
 (0)