Coverage for faststream / redis / parser / binary.py: 98%

1import enum (empty)

2from collections.abc import Sequence (empty)

3from struct import pack, unpack (empty)

4from typing import TYPE_CHECKING, Any, Optional, Union (empty)

5 

6from faststream._internal._compat import json_loads (empty)

7 

8from .message import MessageFormat (empty)

9 

10if TYPE_CHECKING: (empty)

11 from fast_depends.library.serializer import SerializerProto 

12 

13 from faststream._internal.basic_types import SendableMessage 

14 

15 

16class FastStreamMessageVersion(int, enum.Enum): (empty)

17 v1 = 1 (empty)

18 

19 

20class BinaryMessageFormatV1(MessageFormat): (empty)

21 """Message format to encode into binary and parse it.""" 

22 

23 IDENTITY_HEADER = b"\x89BIN\x0d\x0a\x1a\x0a" # to avoid confusion with other formats (empty)

24 

25 @classmethod (empty)

26 def encode( (empty)

27 cls, 

28 *, 

29 message: Union[Sequence["SendableMessage"], "SendableMessage"], 

30 reply_to: str | None, 

31 headers: dict[str, Any] | None, 

32 correlation_id: str, 

33 serializer: Optional["SerializerProto"] = None, 

34 ) -> bytes: 

35 msg = cls.build( (empty)

36 message=message, 

37 reply_to=reply_to, 

38 headers=headers, 

39 correlation_id=correlation_id, 

40 serializer=serializer, 

41 ) 

42 headers_writer = BinaryWriter() (empty)

43 for key, value in msg.headers.items(): (empty)

44 headers_writer.write_string(key) (empty)

45 headers_writer.write_string(value) (empty)

46 

47 headers_len = len(headers_writer.data) (empty)

48 writer = BinaryWriter() (empty)

49 writer.write(cls.IDENTITY_HEADER) (empty)

50 writer.write_short(FastStreamMessageVersion.v1.value) (empty)

51 headers_start = len(writer.data) + 8 (empty)

52 data_start = 2 + headers_start + headers_len (empty)

53 writer.write_int(headers_start) (empty)

54 writer.write_int(data_start) (empty)

55 writer.write_short(len(msg.headers.items())) (empty)

56 writer.write(headers_writer.get_bytes()) (empty)

57 writer.write(msg.data) (empty)

58 return writer.get_bytes() (empty)

59 

60 @classmethod (empty)

61 def parse(cls, data: bytes) -> tuple[bytes, dict[str, Any]]: (empty)

62 headers: dict[str, Any] = {} (empty)

63 final_data: bytes 

64 

65 try: (empty)

66 reader = BinaryReader(data) (empty)

67 magic_header = reader.read_until(len(cls.IDENTITY_HEADER)) (empty)

68 message_version = reader.read_short() (empty)

69 

70 if ( (empty)

71 magic_header == cls.IDENTITY_HEADER 

72 and message_version == FastStreamMessageVersion.v1.value 

73 ): 

74 headers_start = reader.read_int() (empty)

75 data_start = reader.read_int() (empty)

76 reader.shift_offset_to(headers_start) (empty)

77 header_count = reader.read_short() (empty)

78 for _ in range(header_count): (empty)

79 key = reader.read_string() (empty)

80 value = reader.read_string() (empty)

81 headers[key] = value (empty)

82 

83 reader.shift_offset_to(data_start) (empty)

84 final_data = reader.read_bytes() (empty)

85 

86 else: 

87 parsed_data = json_loads(data) (empty)

88 final_data = parsed_data["data"].encode() (empty)

89 headers = parsed_data["headers"] (empty)

90 

91 except Exception: (empty)

92 # Raw Redis message format or legacy JSON envelope 

93 try: (empty)

94 parsed_data = json_loads(data) (empty)

95 final_data = parsed_data["data"].encode() (empty)

96 headers = parsed_data.get("headers", {}) (empty)

97 except Exception: (empty)

98 final_data = data (empty)

99 headers = {} (empty)

100 return final_data, headers (empty)

101 

102 return final_data, headers (empty)

103 

104 

105class BinaryWriter: (empty)

106 def __init__(self) -> None: (empty)

107 self.data = bytearray() (empty)

108 

109 def write(self, data: bytes) -> None: (empty)

110 self.data.extend(data) (empty)

111 

112 def write_short(self, number: int) -> None: (empty)

113 int_bytes = pack(">H", number) (empty)

114 self.write(int_bytes) (empty)

115 

116 def write_int(self, number: int) -> None: (empty)

117 int_bytes = pack(">I", number) (empty)

118 self.write(int_bytes) (empty)

119 

120 def write_string(self, data: str | bytes) -> None: (empty)

121 str_len = len(data) (empty)

122 self.write_short(str_len) (empty)

123 if isinstance(data, bytes): 123 ↛ 124line 123 didn't jump to line 124 because the condition on line 123 was never true(empty)

124 self.write(data) 

125 else: 

126 self.write(data.encode()) (empty)

127 

128 def get_bytes(self) -> bytes: (empty)

129 return bytes(self.data) (empty)

130 

131 

132class BinaryReader: (empty)

133 def __init__(self, data: bytes) -> None: (empty)

134 self.data = data (empty)

135 self.offset = 0 (empty)

136 

137 def read_until(self, offset: int) -> bytes: (empty)

138 data = self.data[self.offset : self.offset + offset] (empty)

139 self.offset += offset (empty)

140 return data (empty)

141 

142 def shift_offset_to(self, offset: int) -> None: (empty)

143 self.offset = offset (empty)

144 

145 def read_short(self) -> int: (empty)

146 data = unpack(">H", self.data[self.offset : self.offset + 2])[0] (empty)

147 self.offset += 2 (empty)

148 return int(data) (empty)

149 

150 def read_int(self) -> int: (empty)

151 data = unpack(">I", self.data[self.offset : self.offset + 4])[0] (empty)

152 self.offset += 4 (empty)

153 return int(data) (empty)

154 

155 def read_string(self) -> str: (empty)

156 str_len = self.read_short() (empty)

157 data = self.data[self.offset : self.offset + str_len] (empty)

158 self.offset += str_len (empty)

159 return data.decode() (empty)

160 

161 def read_bytes(self) -> bytes: (empty)

162 return self.data[self.offset :] (empty)