Skip to content

Commit 1b8bddd

Browse files
shsmsshsms
committed
component_type
Signed-off-by: Sahas Subramanian <[email protected]>
1 parent 2374d6c commit 1b8bddd

File tree

4 files changed

+76
-34
lines changed

4 files changed

+76
-34
lines changed

‎python/frequenz/microgrid_component_graph/__init__.pyi‎

Lines changed: 15 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55

66
"""A graph representation of the electrical components in a microgrid."""
77

8-
from enum import Enum
8+
from collections import abc
99
from typing import AbstractSet, Generic, Protocol, TypeVar
1010

1111
class InvalidGraphError(Exception):
@@ -90,8 +90,8 @@ class ComponentGraph(Generic[ComponentT, ConnectionT]):
9090

9191
def components(
9292
self,
93-
component_category: Enum | None = None,
94-
component_type: Enum | None = None,
93+
filter_by_ids: set[ComponentIdProtocol] | None = None,
94+
filter_by_types: set[type[ComponentProtocol]] | None = None,
9595
) -> set[ComponentT]:
9696
"""Fetch all components in this graph.
9797
@@ -205,28 +205,30 @@ class ComponentGraph(Generic[ComponentT, ConnectionT]):
205205
The grid formula as a string.
206206
"""
207207

208-
def pv_formula(self) -> str:
208+
def pv_formula(self, pv_inverter_ids: abc.Set[ComponentIdProtocol] | None) -> str:
209209
"""Generate the PV formula for this component graph.
210210
211211
Returns:
212212
The PV formula as a string.
213213
"""
214214

215-
def battery_formula(self) -> str:
215+
def battery_formula(self, battery_ids: abc.Set[ComponentIdProtocol] | None) -> str:
216216
"""Generate the battery formula for this component graph.
217217
218218
Returns:
219219
The battery formula as a string.
220220
"""
221221

222-
def chp_formula(self) -> str:
222+
def chp_formula(self, chp_ids: abc.Set[ComponentIdProtocol] | None) -> str:
223223
"""Generate the CHP formula for this component graph.
224224
225225
Returns:
226226
The CHP formula as a string.
227227
"""
228228

229-
def ev_charger_formula(self) -> str:
229+
def ev_charger_formula(
230+
self, ev_charger_ids: abc.Set[ComponentIdProtocol] | None
231+
) -> str:
230232
"""Generate the EV charger formula for this component graph.
231233
232234
Returns:
@@ -240,14 +242,18 @@ class ComponentGraph(Generic[ComponentT, ConnectionT]):
240242
The grid coalesced formula as a string.
241243
"""
242244

243-
def battery_coalesce_formula(self) -> str:
245+
def battery_coalesce_formula(
246+
self, battery_ids: abc.Set[ComponentIdProtocol] | None
247+
) -> str:
244248
"""Generate the battery coalesce formula for this component graph.
245249
246250
Returns:
247251
The battery coalesced formula as a string.
248252
"""
249253

250-
def pv_coalesce_formula(self) -> str:
254+
def pv_coalesce_formula(
255+
self, pv_inverter_ids: abc.Set[ComponentIdProtocol] | None
256+
) -> str:
251257
"""Generate the PV coalesce formula for this component graph.
252258
253259
Returns:

‎src/category.rs‎

Lines changed: 19 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -40,25 +40,35 @@ pub(crate) fn category_from_python_component(
4040
) -> PyResult<cg::ComponentCategory> {
4141
let comp_classes = ComponentClasses::try_new(py)?;
4242

43-
if object.is_instance(&comp_classes.grid_connection_point)? {
43+
if object.is_instance(&comp_classes.grid_connection_point)?
44+
|| object.is(&comp_classes.grid_connection_point)
45+
{
4446
Ok(cg::ComponentCategory::GridConnectionPoint)
45-
} else if object.is_instance(&comp_classes.meter)? {
47+
} else if object.is_instance(&comp_classes.meter)? || object.is(&comp_classes.meter) {
4648
Ok(cg::ComponentCategory::Meter)
47-
} else if object.is_instance(&comp_classes.battery)? {
49+
} else if object.is_instance(&comp_classes.battery)? || object.is(&comp_classes.battery) {
4850
Ok(cg::ComponentCategory::Battery(cg::BatteryType::Unspecified))
49-
} else if object.is_instance(&comp_classes.ev_charger)? {
51+
} else if object.is_instance(&comp_classes.ev_charger)? || object.is(&comp_classes.ev_charger) {
5052
Ok(cg::ComponentCategory::EvCharger(
5153
cg::EvChargerType::Unspecified,
5254
))
53-
} else if object.is_instance(&comp_classes.chp)? {
55+
} else if object.is_instance(&comp_classes.chp)? || object.is(&comp_classes.chp) {
5456
Ok(cg::ComponentCategory::Chp)
55-
} else if object.is_instance(&comp_classes.battery_inverter)? {
57+
} else if object.is_instance(&comp_classes.battery_inverter)?
58+
|| object.is(&comp_classes.battery_inverter)
59+
{
5660
Ok(cg::ComponentCategory::Inverter(cg::InverterType::Battery))
57-
} else if object.is_instance(&comp_classes.solar_inverter)? {
61+
} else if object.is_instance(&comp_classes.solar_inverter)?
62+
|| object.is(&comp_classes.solar_inverter)
63+
{
5864
Ok(cg::ComponentCategory::Inverter(cg::InverterType::Pv))
59-
} else if object.is_instance(&comp_classes.hybrid_inverter)? {
65+
} else if object.is_instance(&comp_classes.hybrid_inverter)?
66+
|| object.is(&comp_classes.hybrid_inverter)
67+
{
6068
Ok(cg::ComponentCategory::Inverter(cg::InverterType::Hybrid))
61-
} else if object.is_instance(&comp_classes.unspecified_component)? {
69+
} else if object.is_instance(&comp_classes.unspecified_component)?
70+
|| object.is(&comp_classes.unspecified_component)
71+
{
6272
Ok(cg::ComponentCategory::Unspecified)
6373
} else {
6474
Err(exceptions::PyValueError::new_err(format!(

‎src/graph.rs‎

Lines changed: 36 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ use std::collections::BTreeSet;
55

66
use crate::{
77
InvalidGraphError,
8-
category::{category_from_python_category_and_type, match_category},
8+
category::{category_from_python_component, match_category},
99
component::Component,
1010
connection::Connection,
1111
utils::extract_int,
@@ -106,27 +106,42 @@ impl ComponentGraph {
106106
cls.into()
107107
}
108108

109-
#[pyo3(signature = (component_category = None, component_type = None))]
109+
#[pyo3(signature = (filter_by_ids=None, filter_by_types=None))]
110110
fn components(
111111
&self,
112-
component_category: Option<Bound<'_, PyAny>>,
113-
component_type: Option<Bound<'_, PyAny>>,
112+
filter_by_ids: Option<Bound<'_, PyAny>>,
113+
filter_by_types: Option<Bound<'_, PyAny>>,
114114
) -> PyResult<Py<PySet>> {
115115
let iter = self.graph.components();
116116

117117
Python::attach(|py| {
118-
if let Some(component_category) = component_category {
119-
let py_category =
120-
category_from_python_category_and_type(component_category, component_type)?;
121-
PySet::new(
122-
py,
123-
iter.filter(|c| match_category(py_category, c.category))
124-
.map(|c| c.object.bind(py)),
125-
)
118+
let components: Vec<_> = if let Some(component_categories) = filter_by_types {
119+
let categories: Vec<cg::ComponentCategory> = component_categories
120+
.try_iter()?
121+
.map(|c| category_from_python_component(py, &c?))
122+
.collect::<PyResult<_>>()?;
123+
124+
iter.filter(|c| categories.iter().any(|x| match_category(*x, c.category)))
125+
.collect()
126126
} else {
127-
PySet::new(py, iter.map(|c| c.object.bind(py)))
128-
}
129-
.map(|s| s.into())
127+
iter.collect()
128+
};
129+
130+
let components: Vec<_> = if let Some(ids) = filter_by_ids {
131+
let ids_set: BTreeSet<u64> = ids
132+
.try_iter()?
133+
.map(|id| extract_int::<u64>(py, id?))
134+
.collect::<PyResult<_>>()?;
135+
136+
components
137+
.into_iter()
138+
.filter(|c| ids_set.contains(&c.component_id))
139+
.collect()
140+
} else {
141+
components
142+
};
143+
144+
PySet::new(py, components.iter().map(|c| c.object.bind(py))).map(|s| s.into())
130145
})
131146
}
132147

@@ -212,6 +227,7 @@ impl ComponentGraph {
212227
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))
213228
}
214229

230+
#[pyo3(signature = (battery_ids=None))]
215231
fn battery_formula(
216232
&self,
217233
py: Python<'_>,
@@ -223,13 +239,15 @@ impl ComponentGraph {
223239
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))
224240
}
225241

242+
#[pyo3(signature = (chp_ids=None))]
226243
fn chp_formula(&self, py: Python<'_>, chp_ids: Option<Bound<'_, PyAny>>) -> PyResult<String> {
227244
self.graph
228245
.chp_formula(extract_ids(py, chp_ids)?)
229246
.map(|f| f.to_string())
230247
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))
231248
}
232249

250+
#[pyo3(signature = (pv_inverter_ids=None))]
233251
fn pv_formula(
234252
&self,
235253
py: Python<'_>,
@@ -241,6 +259,7 @@ impl ComponentGraph {
241259
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))
242260
}
243261

262+
#[pyo3(signature = (ev_charger_ids=None))]
244263
fn ev_charger_formula(
245264
&self,
246265
py: Python<'_>,
@@ -259,6 +278,7 @@ impl ComponentGraph {
259278
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))
260279
}
261280

281+
#[pyo3(signature = (battery_ids=None))]
262282
fn battery_ac_coalesce_formula(
263283
&self,
264284
py: Python<'_>,
@@ -270,6 +290,7 @@ impl ComponentGraph {
270290
.map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(e.to_string()))
271291
}
272292

293+
#[pyo3(signature = (pv_inverter_ids=None))]
273294
fn pv_ac_coalesce_formula(
274295
&self,
275296
py: Python<'_>,

‎tests/test_microgrid_component_graph.py‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,4 +27,9 @@ def test_loading() -> None:
2727
ComponentConnection(source=ComponentId(1), destination=ComponentId(2))
2828
},
2929
)
30-
assert graph is not None
30+
assert graph.components() == {
31+
GridConnectionPoint(
32+
id=ComponentId(1), microgrid_id=MicrogridId(1), rated_fuse_current=100
33+
),
34+
Meter(id=ComponentId(2), microgrid_id=MicrogridId(1)),
35+
}

0 commit comments

Comments
 (0)