Coverage for gramex\handlers\websockethandler.py: 51%

# Coverage for gramex\handlers\websockethandler.py : 51%

## 23 statements 13 run 10 missing 0 excluded 0 partial

Hot-keys on this page

r m x p toggle line displays

j k next/prev highlighted chunk

0 (zero) top of page

1 (one) first highlighted chunk

1 from collections import OrderedDict

2 from six import string_types

3 from six.moves.urllib_parse import urlparse

4 from gramex.transforms import build_transform

5 from .basehandler import BaseWebSocketHandler

6

7

8 class WebSocketHandler(BaseWebSocketHandler):

9 '''

10 Handles WebSockets. It accepts these parameters:

11

12 :arg function open: ``open(handler)`` is called when the connection is opened

13 :arg function on_message: ``on_message(handler, message)`` is called with a

14 string message when the client sends a message.

15 :arg function on_close: ``on_close(handler)`` is called when the websocket is

16 closed.

17 :arg list origins: a domain name or list of domain names. No wildcards

18

19 The handler has a ``.write_message(text)`` method that sends a message back

20 to the client.

21 '''

22 @classmethod

23 def setup(cls, **kwargs):

24 super(WebSocketHandler, cls).setup(**kwargs)

25 override_methods = {

26 'open': ['handler'],

27 'on_message': ['handler', 'message'],

28 'on_close': ['handler'],

29 'on_pong': ['handler', 'data'],

30 'select_subprotocol': ['handler', 'subprotocols'],

31 'get_compression_options': ['handler'],

32 }

33 for method in override_methods:

34 if method in kwargs:

35 setattr(cls, method, build_transform(

36 kwargs[method],

37 vars=OrderedDict((arg, None) for arg in override_methods[method]),

38 filename='url:%s.%s' % (cls.name, method)))

39

40 def check_origin(self, origin):

41 origins = self.kwargs.get('origins', [])

42 if not origins:

43 return True

44 if isinstance(origins, string_types):

45 origins = [origins]

46 domain = urlparse(origin).netloc

47 for allowed_origin in origins:

48 if domain.netloc.endswith(allowed_origin):

49 return True

50 return False