Coverage for src/galaxy_beep_frame/converter/lexer.py: 100%

1"""Lexer class.""" 

2 

3from dataclasses import dataclass, field 

4from typing import TYPE_CHECKING, cast 

5 

6from galaxy_graph import Connection, Graph, Node 

7from galaxy_graph.connection import ConnectionSetup 

8 

9from galaxy_beep_frame.exception.converter import ( 

10 BadEndError, 

11 BadHeaderEndError, 

12 InvalidByteError, 

13 LexerError, 

14 NoCurrentTokenError, 

15 NodeNoneError, 

16 NotFoundError, 

17) 

18 

19from .state import State 

20from .token import Token 

21from .token.token_type import TokenType 

22 

23if TYPE_CHECKING: 

24 from asyncio import StreamReader 

25 from collections.abc import AsyncIterator, Callable, Iterator 

26 from io import Reader 

27 

28 

29@dataclass 

30class Lexer: 

31 """ 

32 Lexer convert a bytes into tokens graph. 

33 

34 Notes 

35 ----- 

36 The graph structure is a hierarchical structure in the graph: 

37 

38 root (id: root) 

39 + frame (has frame) 

40 - header (has header) 

41 - message type (has message type) 

42 - channel number (has channel number) 

43 - message number (has message number) 

44 - continuation indicator (has continuation indicator) 

45 - sequence number (has sequence number) 

46 . answer number (has answer number) 

47 - payload 

48 

49 where: 

50 - + : many 

51 - - : exactly one 

52 - . : one or zero 

53 

54 Frame can be fetched from the graph. They are considered as 

55 completed when in a immutable state. 

56 

57 """ 

58 

59 graph: Graph = field(default_factory=Graph) 

60 

61 frame: Node[Token] | None = None 

62 header: Node[Token] | None = None 

63 root: Node[Token] = field(init=False) 

64 

65 size: int = 0 

66 continuation_indicator: int | None = None 

67 

68 end: bytearray = field(default_factory=bytearray) 

69 

70 current_token: Token | None = None 

71 

72 state: State = State.START 

73 

74 _handlers: dict[State, Callable[[int], Node[Token] | None]] = field( 

75 init=False 

76 ) 

77 

78 def __post_init__(self) -> None: 

79 """Post initialization to save handlers once.""" 

80 self._handlers = { 

81 State.START: self._start, 

82 State.MSG_TYPE: self._msg_type, 

83 State.CHANNEL: self._channel, 

84 State.MESSAGE: self._message, 

85 State.CONTINUATION: self._continuation, 

86 State.SEQUENCE: self._sequence, 

87 State.SIZE: self._size, 

88 State.ANSWER: self._answer, 

89 State.HEADER_END: self._header_end, 

90 State.PAYLOAD: self._payload, 

91 State.END: self._end, 

92 } 

93 

94 self.root = Node( 

95 identifier="root", 

96 data=Token( 

97 TokenType.ROOT, 

98 ), 

99 ) 

100 self.graph.add(self.root) 

101 

102 def _reset_frame(self) -> None: 

103 """Reset property when starting a new frame.""" 

104 

105 self.state = State.START 

106 

107 self.end.clear() 

108 self.size = 0 

109 

110 self.graph = Graph() 

111 self.header = None 

112 self.frame = None 

113 

114 self.current_token = None 

115 

116 def _ensure_current_token(self) -> Token: 

117 """Ensure the current token is not None.""" 

118 if self.current_token is None: 

119 raise NoCurrentTokenError 

120 

121 return self.current_token 

122 

123 @staticmethod 

124 def _ensure_data(node: Node[Token]) -> Token: 

125 """ 

126 Ensure the given node store a Token. 

127 

128 Parameters 

129 ---------- 

130 node: Node[Token] 

131 Node to check. 

132 

133 Returns 

134 ------- 

135 Token 

136 

137 """ 

138 data = node.data 

139 

140 if data is None: 

141 raise NodeNoneError 

142 

143 return data 

144 

145 def _new_token( 

146 self, 

147 token_type: TokenType, 

148 relationship: str, 

149 char: int | None = None, 

150 ) -> None: 

151 """ 

152 Create a new Token and attach it. 

153 

154 The parent is selected implicitly from the parsing context. 

155 

156 Parameters 

157 ---------- 

158 token_type: TokenType: 

159 Type of the new token. 

160 relationship: str 

161 Relationship name between the parent and the new token. 

162 char: int | None 

163 Initial value of the Token. 

164 

165 """ 

166 

167 value = bytearray() 

168 if char is not None: 

169 value.append(char) 

170 

171 if ( 

172 self.header is None 

173 and self.frame is None 

174 and self.state is State.START 

175 ): 

176 # in START, header not intialized 

177 self.frame = Node( 

178 data=Token( 

179 type=TokenType.FRAME, 

180 ) 

181 ) 

182 self.root.connect(Connection("has frame"), self.frame) 

183 self.header = Node( 

184 data=Token( 

185 type=TokenType.HEADER, 

186 ) 

187 ) 

188 self.frame.connect(Connection("has header"), self.header) 

189 parent_node = self.header 

190 elif ( 

191 self.frame is not None 

192 and self.header is not None 

193 and self.state 

194 in { 

195 State.CHANNEL, 

196 State.MESSAGE, 

197 State.CONTINUATION, 

198 State.SEQUENCE, 

199 State.ANSWER, 

200 } 

201 ): 

202 parent_node = self.header 

203 elif ( 

204 self.frame is not None 

205 and self.header is not None 

206 and self.state is State.PAYLOAD 

207 ): 

208 parent_node = self.frame 

209 else: 

210 message = ( 

211 "Cannot create a new token with the current context: " 

212 f"{self.state}" 

213 ) 

214 raise LexerError(message) 

215 

216 self.current_token = Token(type=token_type, value=value) 

217 parent_node.connect( 

218 Connection(relationship), 

219 Node(data=self.current_token), 

220 ) 

221 

222 def _start(self, char: int) -> None: 

223 """Read a byte in START state.""" 

224 

225 if char not in b"ABCDEFGHIJKLMNOPQRSTUVWXYZ": 

226 raise InvalidByteError(char, "an upper case letter") 

227 

228 self._new_token( 

229 TokenType.MESSAGE_TYPE, 

230 "has message type", 

231 char, 

232 ) 

233 self.state = State.MSG_TYPE 

234 

235 def _msg_type(self, char: int) -> None: 

236 """Read a byte in MSG_TYPE state.""" 

237 current_token = self._ensure_current_token() 

238 

239 if char == ord(" "): 

240 current_token.end() 

241 self.state = State.CHANNEL 

242 self._new_token( 

243 TokenType.CHANNEL_NUMBER, 

244 "has channel number", 

245 ) 

246 return 

247 

248 if char not in b"ABCDEFGHIJKLMNOPQRSTUVWXYZ": 

249 raise InvalidByteError(char, "an upper case letter") 

250 

251 current_token.update(char) 

252 

253 def _channel(self, char: int) -> None: 

254 """Read a byte in CHANNEL state.""" 

255 current_token = self._ensure_current_token() 

256 

257 if char == ord(" "): 

258 if len(current_token.value) == 0: 

259 raise NotFoundError(element_name="channel number") 

260 

261 current_token.end() 

262 self.state = State.MESSAGE 

263 self._new_token( 

264 TokenType.MESSAGE_NUMBER, 

265 "has message number", 

266 ) 

267 return 

268 

269 if char not in b"0123456789": 

270 raise InvalidByteError(char, "a digit") 

271 

272 current_token.update(char) 

273 

274 def _message(self, char: int) -> None: 

275 """Read a byte in MESSAGE state.""" 

276 current_token = self._ensure_current_token() 

277 

278 if char == ord(" "): 

279 if len(current_token.value) == 0: 

280 raise NotFoundError(element_name="message number") 

281 

282 current_token.end() 

283 self.state = State.CONTINUATION 

284 self._new_token( 

285 TokenType.CONTINUATION_INDICATOR, 

286 "has continuation indicator", 

287 ) 

288 return 

289 

290 if char not in b"0123456789": 

291 raise InvalidByteError(char, "a digit") 

292 

293 current_token.update(char) 

294 

295 def _continuation(self, char: int) -> None: 

296 """Read a byte in CONTINUATION state.""" 

297 current_token = self._ensure_current_token() 

298 

299 if char == ord(" ") and self.continuation_indicator is not None: 

300 current_token.end() 

301 self.state = State.SEQUENCE 

302 self._new_token( 

303 TokenType.SEQUENCE_NUMBER, 

304 "has sequence number", 

305 ) 

306 self.continuation_indicator = None 

307 return 

308 

309 if char not in b".*": 

310 raise InvalidByteError(char, "in ['*', '.']") 

311 

312 if self.continuation_indicator is not None: 

313 raise InvalidByteError(char, "a space") 

314 

315 self.continuation_indicator = char 

316 current_token.update(char) 

317 

318 def _sequence(self, char: int) -> None: 

319 """Read a byte in SEQUENCE state.""" 

320 current_token = self._ensure_current_token() 

321 

322 if char == ord(" "): 

323 if len(current_token.value) == 0: 

324 raise NotFoundError(element_name="sequence number") 

325 

326 current_token.end() 

327 self.state = State.SIZE 

328 return 

329 

330 if char not in b"0123456789": 

331 raise InvalidByteError(char, "a digit") 

332 

333 current_token.update(char) 

334 

335 def _size(self, char: int) -> None: 

336 """Read a byte in SIZE state.""" 

337 if char == ord("\r"): 

338 self.state = State.HEADER_END 

339 return 

340 

341 if char == ord(" "): 

342 self.state = State.ANSWER 

343 self._new_token( 

344 TokenType.ANSWER_NUMBER, 

345 "has answer number", 

346 ) 

347 return 

348 

349 if char not in b"0123456789": 

350 raise InvalidByteError(char, "a number") 

351 

352 self.size = self.size * 10 + (char - ord("0")) 

353 

354 def _answer(self, char: int) -> None: 

355 """Read a byte in ANSWER state.""" 

356 current_token = self._ensure_current_token() 

357 

358 if char == ord("\r"): 

359 if len(current_token.value) == 0: 

360 raise NotFoundError(element_name="answer number") 

361 

362 current_token.end() 

363 self.state = State.HEADER_END 

364 return 

365 

366 if char not in b"0123456789": 

367 raise InvalidByteError(char, "a digit") 

368 

369 current_token.update(char) 

370 

371 def _header_end(self, char: int) -> None: 

372 """Read a byte in HEADER_END state.""" 

373 

374 # \r is skipped before 

375 if char != ord("\n"): 

376 raise BadHeaderEndError(char) 

377 

378 if self.header is None: 

379 raise NodeNoneError 

380 

381 self._ensure_data(self.header).end() 

382 

383 self.state = State.PAYLOAD 

384 self._new_token(TokenType.PAYLOAD, "has payload") 

385 

386 def _payload(self, char: int) -> None: 

387 """Read a byte in PAYLOAD state.""" 

388 current_token = self._ensure_current_token() 

389 

390 if len(current_token.value) == self.size: 

391 if char != ord("E"): 

392 raise BadEndError(bytes([char])) 

393 current_token.end() 

394 self.end.append(char) 

395 self.state = State.END 

396 return 

397 

398 current_token.update(char) 

399 

400 def _end(self, char: int) -> Node[Token] | None: 

401 """Read a byte in END state.""" 

402 candidate = bytes(self.end) + bytes([char]) 

403 

404 end_token = b"END\r\n" 

405 

406 if candidate == end_token: 

407 if self.frame is None: 

408 raise NodeNoneError 

409 

410 frame = self.frame 

411 self._ensure_data(self.frame).end() 

412 self._reset_frame() 

413 return frame 

414 

415 if not end_token.startswith(candidate): 

416 raise BadEndError(candidate) 

417 

418 self.end.append(char) 

419 

420 return None 

421 

422 def lex_byte(self, char: int) -> Node[Token] | None: 

423 """ 

424 Read a byte with current context. 

425 

426 This function update context in function of the value of this 

427 byte. 

428 """ 

429 

430 return self._handlers[self.state](char) 

431 

432 def feed(self, data: bytes) -> list[Node[Token]]: 

433 """Lex bytes.""" 

434 results: list[Node[Token]] = [] 

435 for byte in data: 

436 result = self.lex_byte(byte) 

437 if result is not None: 

438 results.append(result) 

439 

440 return results 

441 

442 def lex( 

443 self, reader: Reader[bytes], chunk_size: int = 100 

444 ) -> list[Node[Token]]: 

445 """ 

446 Lex on stream of bytes. 

447 

448 Warnings 

449 -------- 

450 This will keep every frames in memory! 

451 

452 """ 

453 results: list[Node[Token]] = [] 

454 while True: 

455 chunk = reader.read(chunk_size) 

456 if not chunk: 

457 return results 

458 results += self.feed(chunk) 

459 

460 def lex_iterator( 

461 self, reader: Reader[bytes], chunk_size: int = 100 

462 ) -> Iterator[Node[Token]]: 

463 """Lex stream of bytes, yield frame node as they complete.""" 

464 while True: 

465 chunk = reader.read(chunk_size) 

466 if not chunk: 

467 return 

468 

469 yield from self.feed(chunk) 

470 

471 async def lex_stream( 

472 self, reader: StreamReader, chunk_size: int = 100 

473 ) -> list[Node[Token]]: 

474 """ 

475 Lex on async stream of bytes. 

476 

477 Warnings 

478 -------- 

479 This will keep every frames in memory! 

480 

481 """ 

482 results: list[Node[Token]] = [] 

483 while True: 

484 chunk: bytes = await reader.read(chunk_size) 

485 if not chunk: 

486 return results 

487 results += self.feed(chunk) 

488 

489 async def lex_stream_iterator( 

490 self, reader: StreamReader, chunk_size: int = 100 

491 ) -> AsyncIterator[Node[Token]]: 

492 """Lex async stream of bytes, yield node as they complete.""" 

493 while True: 

494 chunk: bytes = await reader.read(chunk_size) 

495 if not chunk: 

496 return 

497 

498 for frame_node in self.feed(chunk): 

499 yield frame_node 

500 

501 @staticmethod 

502 def connections_to_human( 

503 connections: set[Connection], indent: int = 1 

504 ) -> str: 

505 """Convert a list of connections to readable str rep.""" 

506 result = "" 

507 for connection in connections: 

508 connection_setup = connection.state.setup 

509 if not isinstance(connection_setup, ConnectionSetup): 

510 message = ( 

511 "Error in the graph, a connection is not setup." 

512 ) 

513 raise LexerError(message) 

514 destination_node = cast( 

515 "Node[Token]", connection_setup.destination 

516 ) 

517 destination = Lexer._ensure_data(destination_node) 

518 result += " " * indent 

519 result += ( 

520 f"{connection.data}: {destination.type} - " 

521 f"`{destination.value.decode()}`\n" 

522 ) 

523 result += Lexer.connections_to_human( 

524 destination_node.outgoing_connections, 

525 indent + 1, 

526 ) 

527 

528 return result 

529 

530 @staticmethod 

531 def node_to_human(node: Node[Token]) -> str: 

532 """Convert a token node to a readable str repr.""" 

533 data = Lexer._ensure_data(node) 

534 result = f"{data.type}\n" 

535 result += "==========\n" 

536 result += Lexer.connections_to_human(node.outgoing_connections) 

537 

538 return result