1313# pylint: disable=no-member
1414from frequenz .api .electricity_trading .v1 import electricity_trading_pb2
1515from frequenz .api .electricity_trading .v1 .electricity_trading_pb2_grpc import (
16+ ElectricityTradingServiceAsyncStub ,
1617 ElectricityTradingServiceStub ,
1718)
1819from frequenz .channels import Receiver
1920from frequenz .client .base .client import BaseApiClient
21+ from frequenz .client .base .exception import ClientNotConnected
2022from frequenz .client .base .streaming import GrpcStreamBroadcaster
2123from frequenz .client .common .pagination import Params
2224from google .protobuf import field_mask_pb2 , struct_pb2
@@ -81,7 +83,7 @@ def validate_decimal_places(value: Decimal, decimal_places: int, name: str) -> N
8183 ) from exc
8284
8385
84- class Client (BaseApiClient [ ElectricityTradingServiceStub ] ):
86+ class Client (BaseApiClient ):
8587 """Electricity trading client."""
8688
8789 _instances : dict [tuple [str , str | None ], "Client" ] = {}
@@ -123,7 +125,11 @@ def __init__(
123125 if not hasattr (
124126 self , "_initialized"
125127 ): # Prevent re-initialization of existing instances
126- super ().__init__ (server_url , ElectricityTradingServiceStub , connect = connect )
128+ super ().__init__ (server_url , connect = connect )
129+ self ._stub = cast (
130+ ElectricityTradingServiceAsyncStub ,
131+ ElectricityTradingServiceStub (self .channel ),
132+ )
127133 self ._initialized = True
128134
129135 self ._gridpool_orders_streams : dict [
@@ -149,6 +155,18 @@ def __init__(
149155
150156 self ._metadata = (("key" , auth_key ),) if auth_key else ()
151157
158+ @property
159+ def stub (self ) -> ElectricityTradingServiceAsyncStub :
160+ """
161+ Get the gRPC stub for the Electricity Trading service.
162+
163+ Returns:
164+ The gRPC stub.
165+ """
166+ if self ._channel is None :
167+ raise ClientNotConnected (server_url = self .server_url , operation = "stub" )
168+ return self ._stub
169+
152170 async def stream_gridpool_orders (
153171 # pylint: disable=too-many-arguments, too-many-positional-arguments
154172 self ,
@@ -192,7 +210,7 @@ async def stream_gridpool_orders(
192210 try :
193211 self ._gridpool_orders_streams [stream_key ] = GrpcStreamBroadcaster (
194212 f"electricity-trading-{ stream_key } " ,
195- lambda : self .stub .ReceiveGridpoolOrdersStream ( # type: ignore
213+ lambda : self .stub .ReceiveGridpoolOrdersStream (
196214 electricity_trading_pb2 .ReceiveGridpoolOrdersStreamRequest (
197215 gridpool_id = gridpool_id ,
198216 filter = gridpool_order_filter .to_pb (),
@@ -251,7 +269,7 @@ async def stream_gridpool_trades(
251269 try :
252270 self ._gridpool_trades_streams [stream_key ] = GrpcStreamBroadcaster (
253271 f"electricity-trading-{ stream_key } " ,
254- lambda : self .stub .ReceiveGridpoolTradesStream ( # type: ignore
272+ lambda : self .stub .ReceiveGridpoolTradesStream (
255273 electricity_trading_pb2 .ReceiveGridpoolTradesStreamRequest (
256274 gridpool_id = gridpool_id ,
257275 filter = gridpool_trade_filter .to_pb (),
@@ -303,7 +321,7 @@ async def stream_public_trades(
303321 self ._public_trades_streams [public_trade_filter ] = (
304322 GrpcStreamBroadcaster (
305323 f"electricity-trading-{ public_trade_filter } " ,
306- lambda : self .stub .ReceivePublicTradesStream ( # type: ignore
324+ lambda : self .stub .ReceivePublicTradesStream (
307325 electricity_trading_pb2 .ReceivePublicTradesStreamRequest (
308326 filter = public_trade_filter .to_pb (),
309327 ),
0 commit comments