Skip to content

Commit 4e38b52

Browse files
committed
Implement extra classes for each target type
..and support the latest CategoryAndType target type. This implements extra classes for each target component case: * List of Ids `TargetIds` * List of categories: `TargetCategories` * List of categories + subtypes: `TargetCategoriesAndTypes` This was easier to do together as the addition of the new type increased the places where it was difficult to find out what the actual target type/ids is/are. Signed-off-by: Mathias L. Baumann <[email protected]>
1 parent d981cd8 commit 4e38b52

File tree

7 files changed

+346
-61
lines changed

7 files changed

+346
-61
lines changed

src/frequenz/client/dispatch/__main__.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,8 @@ def print_dispatch(dispatch: Dispatch) -> None:
8686
# Format the target
8787
if dispatch.target:
8888
if len(dispatch.target) == 1:
89-
target_str: str = str(dispatch.target[0])
89+
(el,) = dispatch.target
90+
target_str: str = str(el)
9091
else:
9192
target_str = ", ".join(str(s) for s in dispatch.target)
9293
else:

src/frequenz/client/dispatch/_cli_types.py

Lines changed: 68 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,15 @@
1212
from tzlocal import get_localzone
1313

1414
from frequenz.client.common.microgrid.components import ComponentCategory
15+
from frequenz.client.dispatch.types import (
16+
BatteryType,
17+
CategoryAndType,
18+
EvChargerType,
19+
InverterType,
20+
TargetCategories,
21+
TargetCategoryAndTypes,
22+
TargetIds,
23+
)
1524

1625
# Disable a false positive from pylint
1726
# pylint: disable=inconsistent-return-statements
@@ -140,7 +149,7 @@ class TargetComponentParamType(click.ParamType):
140149

141150
def convert(
142151
self, value: Any, param: click.Parameter | None, ctx: click.Context | None
143-
) -> list[ComponentCategory] | list[int]:
152+
) -> TargetIds | TargetCategories | TargetCategoryAndTypes:
144153
"""Convert the input value into a list of ComponentCategory or IDs.
145154
146155
Args:
@@ -149,9 +158,13 @@ def convert(
149158
ctx: The Click context object.
150159
151160
Returns:
152-
A list of component ids or component categories.
161+
A list of targets, either as component IDs or component categories.
153162
"""
154-
if isinstance(value, list): # Already a list
163+
if (
164+
isinstance(value, TargetIds)
165+
or isinstance(value, TargetCategories)
166+
or isinstance(value, TargetCategoryAndTypes)
167+
):
155168
return value
156169

157170
values = value.split(",")
@@ -162,20 +175,69 @@ def convert(
162175
error: Exception | None = None
163176
# Attempt to parse component ids
164177
try:
165-
return [int(id) for id in values]
178+
return TargetIds([int(id) for id in values])
166179
except ValueError as e:
167180
error = e
168181

169182
# Attempt to parse as component categories, trim whitespace
170183
try:
171-
return [ComponentCategory[cat.strip().upper()] for cat in values]
184+
return TargetCategories(
185+
[ComponentCategory[cat.strip().upper()] for cat in values]
186+
)
187+
except KeyError as e:
188+
error = e
189+
190+
def valid_short_types() -> (
191+
dict[str, BatteryType | InverterType | EvChargerType]
192+
):
193+
"""All valid types in short form."""
194+
str_to_enum = {}
195+
196+
for enum in BatteryType:
197+
short_name = enum.name.split("_")[-1]
198+
str_to_enum[short_name] = enum
199+
200+
for enum in InverterType:
201+
short_name = enum.name.split("_")[-1]
202+
str_to_enum[short_name] = enum
203+
204+
for enum in EvChargerType:
205+
short_name = enum.name.split("_")[-1]
206+
str_to_enum[short_name] = enum
207+
208+
return str_to_enum
209+
210+
short_types = valid_short_types()
211+
212+
# Attempt to parse as component categories with types
213+
try:
214+
return TargetCategoryAndTypes(
215+
[
216+
CategoryAndType(
217+
ComponentCategory[cat.strip().upper()],
218+
(
219+
short_types[type.strip().upper()]
220+
if type in short_types
221+
else None
222+
),
223+
)
224+
for cat, type in (cat.split(":") for cat in values)
225+
]
226+
)
172227
except KeyError as e:
173228
error = e
174229

175230
self.fail(
176231
f'Invalid component category list or ID list: "{value}".\n'
177232
f'Error: "{error}"\n\n'
178-
"Possible categories: BATTERY, GRID, METER, INVERTER, EV_CHARGER, CHP ",
233+
"Valid formats:\n"
234+
"- 1,2,3 # A list of component IDs\n"
235+
"- METER,INVERTER # A list of component categories\n"
236+
"- BATTERY:NA_ION,INVERTER:SOLAR # A list of component categories with types\n"
237+
"Valid component categories:\n"
238+
f"{', '.join([cat.name for cat in ComponentCategory])}\n"
239+
"Valid component types:\n" # using prepared short_types for better readability
240+
f"{', '.join([f'{cat}:{type}' for cat, type in short_types.items()])}",
179241
param,
180242
ctx,
181243
)

src/frequenz/client/dispatch/_client.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,10 @@
4242
from .types import (
4343
Dispatch,
4444
DispatchEvent,
45+
TargetCategories,
46+
TargetCategoryAndTypes,
4547
TargetComponents,
48+
TargetIds,
4649
_target_components_to_protobuf,
4750
)
4851

@@ -116,7 +119,7 @@ async def list(
116119
self,
117120
microgrid_id: int,
118121
*,
119-
target_components: Iterator[TargetComponents] = iter(()),
122+
target_components: TargetComponents,
120123
start_from: datetime | None = None,
121124
start_to: datetime | None = None,
122125
end_from: datetime | None = None,

src/frequenz/client/dispatch/test/generator.py

Lines changed: 53 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,17 @@
1010

1111
from .._internal_types import rounded_start_time
1212
from ..recurrence import EndCriteria, Frequency, RecurrenceRule, Weekday
13-
from ..types import Dispatch
13+
from ..types import (
14+
BatteryType,
15+
CategoryAndType,
16+
Dispatch,
17+
EvChargerType,
18+
InverterType,
19+
TargetCategories,
20+
TargetCategoryAndTypes,
21+
TargetComponents,
22+
TargetIds,
23+
)
1424

1525

1626
class DispatchGenerator:
@@ -66,6 +76,30 @@ def generate_recurrence_rule(self) -> RecurrenceRule:
6676
],
6777
)
6878

79+
def generate_target_cat_and_type(self) -> CategoryAndType:
80+
"""Generate a random category and type.
81+
82+
Returns:
83+
a random category and type
84+
"""
85+
category = self._rng.choice(list(ComponentCategory)[1:])
86+
type = None
87+
88+
match category:
89+
case ComponentCategory.BATTERY:
90+
type = self._rng.choice(list(BatteryType)[1:])
91+
case ComponentCategory.INVERTER:
92+
type = self._rng.choice(list(InverterType)[1:])
93+
case ComponentCategory.EV_CHARGER:
94+
type = self._rng.choice(list(EvChargerType)[1:])
95+
case _:
96+
type = None
97+
98+
return CategoryAndType(
99+
category=category,
100+
type=type,
101+
)
102+
69103
def generate_dispatch(self) -> Dispatch:
70104
"""Generate a random dispatch instance.
71105
@@ -94,14 +128,24 @@ def generate_dispatch(self) -> Dispatch:
94128
),
95129
target=self._rng.choice( # type: ignore
96130
[
97-
[
98-
self._rng.choice(list(ComponentCategory)[1:])
99-
for _ in range(self._rng.randint(1, 10))
100-
],
101-
[
102-
self._rng.randint(1, 100)
103-
for _ in range(self._rng.randint(1, 10))
104-
],
131+
TargetCategories(
132+
[
133+
self._rng.choice(list(ComponentCategory)[1:])
134+
for _ in range(self._rng.randint(1, 10))
135+
]
136+
),
137+
TargetIds(
138+
[
139+
self._rng.randint(1, 100)
140+
for _ in range(self._rng.randint(1, 10))
141+
]
142+
),
143+
TargetCategoryAndTypes(
144+
[
145+
self.generate_target_cat_and_type()
146+
for _ in range(self._rng.randint(1, 10))
147+
]
148+
),
105149
]
106150
),
107151
active=self._rng.choice([True, False]),

0 commit comments

Comments
 (0)