Coverage for backend/ahuora-builder/src/ahuora_builder/port_manager.py: 81%

1from ahuora_builder_types.unit_model_schema import PortSchema 

2from pyomo.network import Port 

3from ahuora_builder_types import PortId 

4 

5 

6class PortManager: 

7 

8 def __init__(self) -> None: 

9 self._map: dict[PortId, Port] = {} 

10 self._schema_map: dict[PortId, PortSchema] = {} 

11 

12 

13 def register_port(self, schema: PortSchema, port: Port) -> None: 

14 """Registers a port with the port map, so that arcs can connect to it by id""" 

15 id = schema.id 

16 self._map[id] = port 

17 self._schema_map[id] = schema 

18 

19 def get_port_schema(self, id: PortId) -> PortSchema: 

20 try: 

21 return self._schema_map[id] 

22 except KeyError: 

23 raise KeyError(f"Port schema with id `{id}` not found") 

24 

25 

26 def get_port(self, id: PortId) -> Port: 

27 try: 

28 return self._map[id] 

29 except KeyError: 

30 raise KeyError(f"Port with id `{id}` not found")