1+ import asyncio
12import re
3+ from functools import wraps
4+ from typing import Dict , Union , Tuple , List , Optional , Any , TypeVar , Callable , cast
25
36import websockets
47
1114KPageSize = 65536
1215
1316
14- def calculate_offset (aIndex ) :
17+ def calculate_offset (aIndex : int ) -> int :
1518 offset = 0
1619
1720 if (aIndex > vlxDevConstants .RANGE_START_g_cyclone_general_info ) and (
@@ -125,6 +128,28 @@ def to_kelvin(value: float) -> int:
125128 return int (round (value * 10 ) * 10 + 27315 )
126129
127130
131+ FuncT = TypeVar ("FuncT" , bound = Callable [..., Any ])
132+
133+
134+ def _websocket_exception_handler (request_fn : FuncT ) -> FuncT :
135+ @wraps (request_fn )
136+ async def wrapped (* args : Any , ** kwargs : Any ) -> Any :
137+ try :
138+ return await request_fn (* args , ** kwargs )
139+ except websockets .InvalidHandshake as e :
140+ raise ValloxWebsocketException ("Websocket handshake failed" ) from e
141+ except websockets .InvalidURI as e :
142+ raise ValloxWebsocketException ("Websocket invalid URI" ) from e
143+ except websockets .PayloadTooBig as e :
144+ raise ValloxWebsocketException ("Websocket payload too big" ) from e
145+ except websockets .InvalidState as e :
146+ raise ValloxWebsocketException ("Websocket invalid state" ) from e
147+ except websockets .WebSocketProtocolError as e :
148+ raise ValloxWebsocketException ("Websocket protocol error" ) from e
149+
150+ return cast (FuncT , wrapped )
151+
152+
128153class Client :
129154 SETTABLE_INT_VALS = {
130155 re .compile ("^A_CYC_STATE$" ),
@@ -135,18 +160,12 @@ class Client:
135160 re .compile ("^A_CYC_(?:EXTR|SUPP)_FAN_BALANCE_BASE$" ),
136161 }
137162
138- _settable_addresses = None
163+ _settable_addresses : Dict [ int , type ]
139164
140- def get_settable_addresses (self ):
141- if not self ._settable_addresses :
142- self ._settable_addresses = dict (
143- (v , int )
144- for k , v in vlxDevConstants .__dict__ .items ()
145- if any (r .match (k ) for r in self .SETTABLE_INT_VALS )
146- )
165+ def get_settable_addresses (self ) -> Dict [int , type ]:
147166 return self ._settable_addresses
148167
149- def set_settable_address (self , address , var_type ) :
168+ def set_settable_address (self , address : Union [ int , str ], var_type : type ) -> None :
150169 if var_type not in int , float ]:
151170 raise AttributeError ("Only float or int type are supported" )
152171
@@ -165,16 +184,23 @@ def set_settable_address(self, address, var_type):
165184 "Unable to add address '%s' to settable list" % str (address )
166185 )
167186
168- def __init__ (self , ip_address ):
187+ def __init__ (self , ip_address : str ):
169188 self .ip_address = ip_address
170189
171- def _decode_pair (self , key , value ):
190+ self ._settable_addresses = dict (
191+ (v , int )
192+ for k , v in vlxDevConstants .__dict__ .items ()
193+ if any (r .match (k ) for r in self .SETTABLE_INT_VALS )
194+ )
195+
196+ def _decode_pair (self , key : str , value : Union [int ,str ]) -> Tuple [int , Union [int , float ]]:
172197 try :
173198 address = int (getattr (vlxDevConstants , key , key ))
174199 except ValueError :
175200 raise AttributeError ("%s setting does not exist" % key )
176201 if "_TEMP_" in key :
177202 value = to_kelvin (float (value ))
203+ raw_value : Union [int , float ]
178204 try :
179205 raw_value = int (value )
180206 except ValueError :
@@ -185,42 +211,35 @@ def _decode_pair(self, key, value):
185211 required_type = addresses [address ]
186212 except KeyError :
187213 raise AttributeError ("%s setting is not settable" % key )
214+
188215 assert type (raw_value ) == required_type , (
189216 "%s(%d) key needs to be an %s, but %s passed"
190217 % (key , address , required_type .__name__ , type (raw_value ).__name__ )
191218 )
192219
193220 return address , raw_value
194221
195- async def _websocket_request (self , payload , read_packets = 1 ):
196- try :
197- async with websockets .connect ("ws://%s/" % self .ip_address ) as ws :
198- await ws .send (payload )
199- results = []
200- for i in range (0 , read_packets ):
201- r = await ws .recv ()
202-
203- results .append (r )
204- return results [0 ] if read_packets == 1 else results
205- except websockets .InvalidHandshake as e :
206- raise ValloxWebsocketException ("Websocket handshake failed" ) from e
207- except websockets .InvalidURI as e :
208- raise ValloxWebsocketException ("Websocket invalid URI" ) from e
209- except websockets .PayloadTooBig as e :
210- raise ValloxWebsocketException ("Websocket payload too big" ) from e
211- except websockets .InvalidState as e :
212- raise ValloxWebsocketException ("Websocket invalid state" ) from e
213- except websockets .WebSocketProtocolError as e :
214- raise ValloxWebsocketException ("Websocket protocol error" ) from e
222+ @_websocket_exception_handler
223+ async def _websocket_request (self , payload : bytes ) -> bytes :
224+ async with websockets .connect ("ws://%s/" % self .ip_address ) as ws :
225+ await ws .send (payload )
226+ r : bytes = await ws .recv ()
227+ return r
228+
229+ @_websocket_exception_handler
230+ async def _websocket_request_multiple (self , payload : bytes , read_packets : int ) -> List [bytes ]:
231+ async with websockets .connect ("ws://%s/" % self .ip_address ) as ws :
232+ await ws .send (payload )
233+ return await asyncio .gather (* [ws .recv () for _ in range (0 , read_packets )])
215234
216- async def fetch_metrics (self , metric_keys = None ):
235+ async def fetch_metrics (self , metric_keys : Optional [ List [ str ]] = None ) -> Dict [ str , Union [ int , float ]] :
217236 metrics = {}
218237 payload = ReadTableRequest .build ({})
219238 result = await self ._websocket_request (payload )
220239
221240 data = ReadTableResponse .parse (result )
222241
223- if not metric_keys :
242+ if metric_keys is None :
224243 metric_keys = vlxDevConstants .__dict__ .keys ()
225244
226245 for key in metric_keys :
@@ -233,9 +252,9 @@ async def fetch_metrics(self, metric_keys=None):
233252
234253 return metrics
235254
236- async def fetch_raw_logs (self ):
255+ async def fetch_raw_logs (self ) -> List [ List [ Dict [ str , Union [ int , float ]]]] :
237256 payload = LogReadRequest .build ({})
238- result = await self ._websocket_request (payload , read_packets = 2 )
257+ result = await self ._websocket_request_multiple (payload , read_packets = 2 )
239258 page_count = LogReadResponse1 .parse (result [0 ]).fields .value .pages
240259
241260 expected_total_len = KPageSize * page_count
@@ -269,10 +288,10 @@ async def fetch_raw_logs(self):
269288
270289 return series
271290
272- async def fetch_metric (self , metric_key ) :
291+ async def fetch_metric (self , metric_key : str ) -> Optional [ Union [ int , float ]] :
273292 return (await self .fetch_metrics ([metric_key ])).get (metric_key , None )
274293
275- async def set_values (self , dict_ ) :
294+ async def set_values (self , dict_ : Dict [ str , Union [ int , str ]]) -> bool :
276295 items = []
277296 for key , value in dict_ .items ():
278297 address , raw_value = self ._decode_pair (key , value )
0 commit comments