3636from frequenz .client .base .exception import ClientNotConnected
3737from frequenz .client .base .retry import LinearBackoff
3838from frequenz .client .base .streaming import GrpcStreamBroadcaster
39+ from frequenz .client .common .microgrid import MicrogridId
3940
4041from ._internal_types import DispatchCreateRequest
4142from .recurrence import RecurrenceRule
4243from .types import (
4344 Dispatch ,
4445 DispatchEvent ,
46+ DispatchId ,
4547 TargetComponents ,
4648 _target_components_to_protobuf ,
4749)
@@ -83,7 +85,8 @@ def __init__(
8385 )
8486 self ._metadata = (("key" , key ),)
8587 self ._streams : dict [
86- int , GrpcStreamBroadcaster [StreamMicrogridDispatchesResponse , DispatchEvent ]
88+ MicrogridId ,
89+ GrpcStreamBroadcaster [StreamMicrogridDispatchesResponse , DispatchEvent ],
8790 ] = {}
8891 """A dictionary of streamers, keyed by microgrid_id."""
8992
@@ -114,7 +117,7 @@ def stub(self) -> dispatch_pb2_grpc.MicrogridDispatchServiceAsyncStub:
114117 # pylint: disable=too-many-arguments, too-many-locals
115118 async def list (
116119 self ,
117- microgrid_id : int ,
120+ microgrid_id : MicrogridId ,
118121 * ,
119122 target_components : Iterator [TargetComponents ] = iter (()),
120123 start_from : datetime | None = None ,
@@ -138,7 +141,7 @@ async def list(
138141 key="key",
139142 server_url="grpc://dispatch.url.goes.here.example.com"
140143 )
141- async for page in client.list(microgrid_id=1 ):
144+ async for page in client.list(microgrid_id=MicrogridId(1) ):
142145 for dispatch in page:
143146 print(dispatch)
144147 ```
@@ -185,7 +188,7 @@ def to_interval(
185188 )
186189
187190 request = ListMicrogridDispatchesRequest (
188- microgrid_id = microgrid_id ,
191+ microgrid_id = int ( microgrid_id ) ,
189192 filter = filters ,
190193 pagination_params = (
191194 PaginationParams (page_size = page_size ) if page_size else None
@@ -211,7 +214,7 @@ def to_interval(
211214 else :
212215 break
213216
214- def stream (self , microgrid_id : int ) -> channels .Receiver [DispatchEvent ]:
217+ def stream (self , microgrid_id : MicrogridId ) -> channels .Receiver [DispatchEvent ]:
215218 """Receive a stream of dispatch events.
216219
217220 This function returns a receiver channel that can be used to receive
@@ -238,15 +241,15 @@ def stream(self, microgrid_id: int) -> channels.Receiver[DispatchEvent]:
238241 return self ._get_stream (microgrid_id ).new_receiver ()
239242
240243 def _get_stream (
241- self , microgrid_id : int
244+ self , microgrid_id : MicrogridId
242245 ) -> GrpcStreamBroadcaster [StreamMicrogridDispatchesResponse , DispatchEvent ]:
243246 """Get an instance to the streaming helper."""
244247 broadcaster = self ._streams .get (microgrid_id )
245248 if broadcaster is not None and not broadcaster .is_running :
246249 del self ._streams [microgrid_id ]
247250 broadcaster = None
248251 if broadcaster is None :
249- request = StreamMicrogridDispatchesRequest (microgrid_id = microgrid_id )
252+ request = StreamMicrogridDispatchesRequest (microgrid_id = int ( microgrid_id ) )
250253 broadcaster = GrpcStreamBroadcaster (
251254 stream_name = "StreamMicrogridDispatches" ,
252255 stream_method = lambda : cast (
@@ -266,7 +269,7 @@ def _get_stream(
266269
267270 async def create ( # pylint: disable=too-many-positional-arguments
268271 self ,
269- microgrid_id : int ,
272+ microgrid_id : MicrogridId ,
270273 type : str , # pylint: disable=redefined-builtin
271274 start_time : datetime | Literal ["NOW" ],
272275 duration : timedelta | None ,
@@ -334,8 +337,8 @@ async def create( # pylint: disable=too-many-positional-arguments
334337 async def update (
335338 self ,
336339 * ,
337- microgrid_id : int ,
338- dispatch_id : int ,
340+ microgrid_id : MicrogridId ,
341+ dispatch_id : DispatchId ,
339342 new_fields : dict [str , Any ],
340343 ) -> Dispatch :
341344 """Update a dispatch.
@@ -359,7 +362,7 @@ async def update(
359362 ValueError: If updating `type` or `dry_run`.
360363 """
361364 msg = UpdateMicrogridDispatchRequest (
362- dispatch_id = dispatch_id , microgrid_id = microgrid_id
365+ dispatch_id = int ( dispatch_id ) , microgrid_id = int ( microgrid_id )
363366 )
364367
365368 for key , val in new_fields .items ():
@@ -423,7 +426,9 @@ async def update(
423426
424427 return Dispatch .from_protobuf (response .dispatch )
425428
426- async def get (self , * , microgrid_id : int , dispatch_id : int ) -> Dispatch :
429+ async def get (
430+ self , * , microgrid_id : MicrogridId , dispatch_id : DispatchId
431+ ) -> Dispatch :
427432 """Get a dispatch.
428433
429434 Args:
@@ -434,7 +439,7 @@ async def get(self, *, microgrid_id: int, dispatch_id: int) -> Dispatch:
434439 Dispatch: The dispatch.
435440 """
436441 request = GetMicrogridDispatchRequest (
437- dispatch_id = dispatch_id , microgrid_id = microgrid_id
442+ dispatch_id = int ( dispatch_id ) , microgrid_id = int ( microgrid_id )
438443 )
439444 response = await cast (
440445 Awaitable [GetMicrogridDispatchResponse ],
@@ -444,15 +449,17 @@ async def get(self, *, microgrid_id: int, dispatch_id: int) -> Dispatch:
444449 )
445450 return Dispatch .from_protobuf (response .dispatch )
446451
447- async def delete (self , * , microgrid_id : int , dispatch_id : int ) -> None :
452+ async def delete (
453+ self , * , microgrid_id : MicrogridId , dispatch_id : DispatchId
454+ ) -> None :
448455 """Delete a dispatch.
449456
450457 Args:
451458 microgrid_id: The microgrid_id to delete the dispatch for.
452459 dispatch_id: The dispatch_id to delete.
453460 """
454461 request = DeleteMicrogridDispatchRequest (
455- dispatch_id = dispatch_id , microgrid_id = microgrid_id
462+ dispatch_id = int ( dispatch_id ) , microgrid_id = int ( microgrid_id )
456463 )
457464 await cast (
458465 Awaitable [None ],
0 commit comments