19
19
from websocket import WebSocket , ABNF , enableTrace
20
20
import six
21
21
import ssl
22
- from six .moves .urllib .parse import urlencode
23
- from six .moves .urllib .parse import quote_plus
22
+ from six .moves .urllib .parse import urlencode , quote_plus , urlparse , urlunparse
24
23
25
24
STDIN_CHANNEL = 0
26
25
STDOUT_CHANNEL = 1
@@ -203,18 +202,21 @@ def close(self, **kwargs):
203
202
WSResponse = collections .namedtuple ('WSResponse' , ['data' ])
204
203
205
204
205
+ def get_websocket_url (url ):
206
+ parsed_url = urlparse (url )
207
+ parts = list (parsed_url )
208
+ if parsed_url .scheme == 'http' :
209
+ parts [0 ] = 'ws'
210
+ elif parsed_url .scheme == 'https' :
211
+ parts [0 ] = 'wss'
212
+ return urlunparse (parts )
213
+
214
+
206
215
def websocket_call (configuration , url , query_params , _request_timeout ,
207
216
_preload_content , headers ):
208
217
"""An internal function to be called in api-client when a websocket
209
218
connection is required."""
210
219
211
- # switch protocols from http to websocket
212
- url = url .replace ('http://' , 'ws://' )
213
- url = url .replace ('https://' , 'wss://' )
214
-
215
- # patch extra /
216
- url = url .replace ('//api' , '/api' )
217
-
218
220
# Extract the command from the list of tuples
219
221
commands = None
220
222
for key , value in query_params :
@@ -238,7 +240,7 @@ def websocket_call(configuration, url, query_params, _request_timeout,
238
240
url += '&command=' + quote_plus (commands )
239
241
240
242
try :
241
- client = WSClient (configuration , url , headers )
243
+ client = WSClient (configuration , get_websocket_url ( url ) , headers )
242
244
if not _preload_content :
243
245
return client
244
246
client .run_forever (timeout = _request_timeout )
0 commit comments