# Coverage for endpoint.py : 72%

## 123 statements 96 run 27 missing 2 excluded 10 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# -*- coding: utf-8 -*-

2

3# -- stdlib --

4 from enum import IntEnum

5 from typing import Any, Iterator, Sequence, Tuple, List

6 import logging

7 import zlib

8

9# -- third party --

10 from gevent import socket

11 from gevent.lock import RLock

12 from gevent.timeout import Timeout

13 import msgpack

14

15# -- own --

16 import wire

17

18

19# -- code --

20 log = logging.getLogger("Endpoint")

21 log.setLevel(logging.ERROR)

22

23

24 class EndpointDied(Exception):

25 pass

26

27

28 class DecodeError(Exception):

29 pass

30

31

32 class Format(IntEnum):

33 Packed = 1

34 BulkCompressed = 2

35

36

37 class Endpoint(object):

38

39 def __init__(self, sock, address):

40 sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)

41 sock.read = sock.recv

42 sock.write = sock.sendall

43

44 self.sock = sock

45 self.unpacker = msgpack.Unpacker(sock, raw=False, strict_map_key=False)

46 self.writelock = RLock()

47 self.address = address

48 self.link_state = 'connected' # or disconnected

49

50 self.active = False

51

52 def __repr__(self):

53 return '%s:%s:%s' % (

54 self.__class__.__name__,

55 self.address[0],

56 self.address[1],

57 )

58

59 @staticmethod

60 def encode(p: wire.Message) -> bytes:

61 return msgpack.packb([Format.Packed, p.encode()], use_bin_type=True)

62

63 @staticmethod

64 def encode_bulk(pl: Sequence[wire.Message]) -> bytes:

65 data = msgpack.packb([i.encode() for i in pl], use_bin_type=True)

66 return msgpack.packb([Format.BulkCompressed, zlib.compress(data)], use_bin_type=True)

67

68 def write(self, p) -> None:

69 log.debug("%s SEND>> %s", repr(self), p)

70 self.raw_write(self.encode(p))

71

72 def write_bulk(self, pl: Sequence[wire.Message]) -> None:

73 for p in pl:

74 log.debug("%s SEND>> %s", repr(self), p)

75 self.raw_write(self.encode_bulk(pl))

76

77 def raw_write(self, s: bytes) -> None:

78 if self.link_state == 'connected': 78 ↛ exitline 78 didn't return from function 'raw_write', because the condition on line 78 was never false

79 try:

80 with self.writelock:

81 self.sock.sendall(s)

82 self.active = True

83 except IOError:

84 log.exception('Error raw_write, closing Endpoint')

85 self.close()

86

87 def close(self):

88 if not self.link_state == 'disconnected': 88 ↛ exitline 88 didn't return from function 'close', because the condition on line 88 was never false

89 self.link_state = 'disconnected'

90 self.sock.close()

91 self.active = False

92

93 @staticmethod

94 def _decode_packet(p: Any) -> Tuple[Format, Any]:

95 try:

96 if not (isinstance(p, (list, tuple)) and len(p) == 2): 96 ↛ 97line 96 didn't jump to line 97, because the condition on line 96 was never true

97 raise DecodeError

98

99 fmt, data = p

100 fmt = Format(fmt)

101 if fmt == Format.Packed:

102 return fmt, data

103 elif fmt == Format.BulkCompressed: 103 ↛ 111line 103 didn't jump to line 111, because the condition on line 103 was never false

104 try:

105 inflated = zlib.decompress(data)

106 except Exception:

107 raise DecodeError

108

109 return fmt, msgpack.unpackb(inflated, raw=False, ext_hook=lambda code, data: None) 109 ↛ exitline 109 didn't run the lambda on line 109

110 else:

111 raise DecodeError

112 except (ValueError, msgpack.UnpackValueError):

113 raise DecodeError

114

115 @staticmethod

116 def decode_bytes(s: bytes) -> List[wire.Message]:

117 p = msgpack.unpackb(s, raw=False)

118 fmt, data = Endpoint._decode_packet(p)

119 if fmt == Format.Packed:

120 msg = wire.Message.decode(data)

121 if not msg:

122 raise DecodeError

123 return [msg]

124 elif fmt == Format.BulkCompressed:

125 return data

126

127 assert False, 'WTF'

128

129 def messages(self, timeout=90) -> Iterator[wire.Message]:

130 if self.link_state != 'connected': 130 ↛ 131line 130 didn't jump to line 131, because the condition on line 130 was never true

131 raise EndpointDied

132

133 unpacker = self.unpacker

134 _NONE = object()

135 while True:

136 try:

137 v = _NONE

138 with Timeout(timeout, False):

139 try:

140 v = next(unpacker)

141 self.active = True

142 except StopIteration:

143 pass

144

145 if v is _NONE:

146 break

147

148 fmt, data = self._decode_packet(v)

149 if fmt == Format.Packed:

150 msg = wire.Message.decode(data)

151 if msg: 151 ↛ 136line 151 didn't jump to line 136, because the condition on line 151 was never false

152 log.debug("%s <<RECV %r", repr(self), msg)

153 yield msg

154 elif fmt == Format.BulkCompressed: 154 ↛ 136line 154 didn't jump to line 136, because the condition on line 154 was never false

155 for d in data:

156 msg = wire.Message.decode(d)

157 if msg: 157 ↛ 155line 157 didn't jump to line 155, because the condition on line 157 was never false

158 log.debug("%s <<RECV %r", repr(self), msg)

159 yield msg

160

161 except DecodeError: 161 ↛ 162line 161 didn't jump to line 162, because the exception caught by line 161 didn't happen

162 log.info('DecodeError')

163 break

164

165 except msgpack.UnpackValueError:

166 log.exception('UnpackValueError')

167 break

168

169 self.close()

170 raise EndpointDied