1+ from enum import Enum
12from threading import Thread , Lock
23from websocket import WebSocketApp , setdefaulttimeout , ABNF
34from msgpack import packb , unpackb
45from ssl import CERT_NONE
56
67
8+ class ConnectionState (Enum ):
9+ NONE = 0
10+ CONNECTING = 1
11+ CONNECTED = 2
12+ FAILED = 3
713class WSConnector :
814
915 class REID :
@@ -26,7 +32,7 @@ def __init__(self, username: str, token: str, address: str, on_msg=None, ignore_
2632 self .ws = None
2733 self .lock = Lock ()
2834 self .reid = self .REID ()
29- self .running = False
35+ self .__connection_state = ConnectionState . NONE
3036 self .ignore_ssl_cert = ignore_ssl_cert
3137 setdefaulttimeout (60 )
3238
@@ -40,26 +46,32 @@ def start(self):
4046 on_message = None if self .on_msg is None else self ._handle_msg ,
4147 on_open = self ._ready , on_error = self ._fail )
4248 self .lock .acquire () # wait for connection to be established
49+ self .__connection_state = ConnectionState .CONNECTING
4350 kwargs = {"sslopt" : {"cert_reqs" : CERT_NONE }} if self .ignore_ssl_cert else None
4451 Thread (target = self .ws .run_forever , kwargs = kwargs ).start ()
4552
4653 def _fail (self , ws , err ):
54+ self .__connection_state = ConnectionState .FAILED
4755 self .lock .release ()
4856 raise err
4957
5058 def stop (self ):
5159 if self .ws is not None :
5260 with self .lock :
5361 print ("Closing the connection." )
54- self .running = False
62+ self .__connection_state = ConnectionState . NONE
5563 self .ws .close ()
5664 self .ws = None
5765
5866 def _ready (self , ws ):
5967 print (f"Connected to { self .address } ." )
60- self .running = True
68+ self .__connection_state = ConnectionState . CONNECTED
6169 self .lock .release ()
6270
71+ @property
72+ def connection_state (self ):
73+ return self .__connection_state
74+
6375 def _handle_msg (self , ws , msg ):
6476 if isinstance (msg , bytes ):
6577 msg = unpackb (msg )
0 commit comments