1c"""source-before-namespace 

2#include <cstring> 

3""" 

4 

5class BsonError(Error): 

6 message: string 

7 

8class _Reader: 

9 _data: bytes 

10 _pos: i64 

11 

12 func __init__(self, data: bytes): 

13 self._data = data 

14 self._pos = 0 

15 

16 func available(self) -> bool: 

17 return self._pos < self._data.length() 

18 

19 func read_u8(self) -> u8: 

20 if not self.available(): 

21 raise BsonError("No more data available.") 

22 

23 value = self._data[self._pos] 

24 self._pos += 1 

25 

26 return value 

27 

28 func read_u16(self) -> u16: 

29 return u16(self.read_u8()) | (u16(self.read_u8()) << 8) 

30 

31 func read_u32(self) -> u32: 

32 return u32(self.read_u16()) | (u32(self.read_u16()) << 16) 

33 

34 func read_u64(self) -> u64: 

35 return u64(self.read_u32()) | (u64(self.read_u32()) << 32) 

36 

37 func read_i32(self) -> i32: 

38 return i32(self.read_u32()) 

39 

40 func read_i64(self) -> i64: 

41 return i64(self.read_u64()) 

42 

43 func read_f64(self) -> f64: 

44 value = 0.0 

45 ivalue = self.read_u64() 

46 c"memcpy(&value, &ivalue, sizeof(value));" 

47 

48 return value 

49 

50 func read_bytes(self, count: i64) -> bytes: 

51 data = b"" 

52 

53 for _ in range(count): 

54 data += self.read_u8() 

55 

56 return data 

57 

58class _Writer: 

59 _data: bytes 

60 

61 func __init__(self): 

62 self._data = b"" 

63 

64 func data(self) -> bytes: 

65 return self._data 

66 

67 func write_i32_at(self, offset: i64, value: i32): 

68 self._data[offset + 0] = u8(value) 

69 self._data[offset + 1] = u8(value >> 8) 

70 self._data[offset + 2] = u8(value >> 16) 

71 self._data[offset + 3] = u8(value >> 24) 

72 

73 func write_u8(self, value: u8): 

74 self._data += value 

75 

76 func write_u16(self, value: u16): 

77 self.write_u8(u8(value)) 

78 self.write_u8(u8(value >> 8)) 

79 

80 func write_u32(self, value: u32): 

81 self.write_u16(u16(value)) 

82 self.write_u16(u16(value >> 16)) 

83 

84 func write_u64(self, value: u64): 

85 self.write_u32(u32(value)) 

86 self.write_u32(u32(value >> 32)) 

87 

88 func write_i32(self, value: i32): 

89 self.write_u32(u32(value)) 

90 

91 func write_i64(self, value: i64): 

92 self.write_u64(u64(value)) 

93 

94 func write_f64(self, value: f64): 

95 ivalue = 0 

96 c"memcpy(&ivalue, &value, sizeof(ivalue));" 

97 self.write_i64(ivalue) 

98 

99 func write_bytes(self, value: bytes): 

100 self._data += value 

101 

102trait Element: 

103 """An element in a BSON document. 

104 

105 """ 

106 

107 func double(self) -> f64: 

108 raise BsonError("Not a double.") 

109 

110 func string(self) -> string: 

111 raise BsonError("Not a string.") 

112 

113 func document(self) -> [(string, Element)]: 

114 raise BsonError("Not a document.") 

115 

116 func get(self, name: string) -> Element: 

117 raise BsonError("Not a document.") 

118 

119 func array(self) -> [Element]: 

120 raise BsonError("Not an array.") 

121 

122 func at(self, index: i64) -> Element: 

123 raise BsonError("Not an array.") 

124 

125 func binary(self) -> (u8, bytes): 

126 raise BsonError("Not a binary.") 

127 

128 func boolean(self) -> bool: 

129 raise BsonError("Not a boolean.") 

130 

131 func int32(self) -> i32: 

132 raise BsonError("Not an int32.") 

133 

134 func int64(self) -> i64: 

135 raise BsonError("Not an int64.") 

136 

137class Double(Element): 

138 """A double. 

139 

140 """ 

141 

142 value: f64 

143 

144 func double(self) -> f64: 

145 return self.value 

146 

147class String(Element): 

148 """A string. 

149 

150 """ 

151 

152 value: string 

153 

154 func string(self) -> string: 

155 return self.value 

156 

157class Document(Element): 

158 """A document. 

159 

160 """ 

161 

162 elements: [(string, Element)] 

163 

164 func __init__(self, elements: [(string, Element)] = []): 

165 self.elements = elements 

166 

167 func document(self) -> [(string, Element)]: 

168 return self.elements 

169 

170 func get(self, name: string) -> Element: 

171 for key, element in self.elements: 

172 if key == name: 

173 return element 

174 

175 raise BsonError(f"Element '{name}' not found.") 

176 

177class Array(Element): 

178 """An array. 

179 

180 """ 

181 

182 doc: Document 

183 

184 func __init__(self, elements: [Element] = []): 

185 self.doc = Document() 

186 

187 for i, element in enumerate(elements): 

188 self.doc.elements.append((str(i), element)) 

189 

190 func append(self, element: Element): 

191 index = str(self.doc.elements.length()) 

192 self.doc.elements.append((index, element)) 

193 

194 func array(self) -> [Element]: 

195 return [element for _, element in self.doc.elements] 

196 

197 func at(self, index: i64) -> Element: 

198 return self.doc.elements[index][1] 

199 

200class Binary(Element): 

201 """A binary blob. 

202 

203 """ 

204 

205 value: (u8, bytes) 

206 

207 func binary(self) -> (u8, bytes): 

208 return self.value 

209 

210class ObjectId(Element): 

211 """An object id. 

212 

213 """ 

214 

215 value: bytes 

216 

217 func object_id(self) -> bytes: 

218 return self.value 

219 

220class Boolean(Element): 

221 """A boolean. 

222 

223 """ 

224 

225 value: bool 

226 

227 func boolean(self) -> bool: 

228 return self.value 

229 

230class Int32(Element): 

231 """An i32. 

232 

233 """ 

234 

235 value: i32 

236 

237 func int32(self) -> i32: 

238 return self.value 

239 

240class Int64(Element): 

241 """An i64. 

242 

243 """ 

244 

245 value: i64 

246 

247 func int64(self) -> i64: 

248 return self.value 

249 

250func _decode_cstring(reader: _Reader) -> string: 

251 data = b"" 

252 

253 while True: 

254 value = reader.read_u8() 

255 

256 if value == 0: 

257 break 

258 

259 data += value 

260 

261 return string(data) 

262 

263func _decode_double(reader: _Reader) -> Double: 

264 return Double(reader.read_f64()) 

265 

266func _decode_string(reader: _Reader) -> String: 

267 length = reader.read_i32() 

268 element = String(string(reader.read_bytes(i64(length - 1)))) 

269 reader.read_u8() 

270 

271 return element 

272 

273func _decode_document(reader: _Reader) -> Document: 

274 reader.read_i32() 

275 document = Document() 

276 

277 while True: 

278 kind = reader.read_u8() 

279 

280 if kind == 0: 

281 break 

282 

283 name = _decode_cstring(reader) 

284 element: Element? = None 

285 

286 match kind: 

287 case 1: 

288 element = _decode_double(reader) 

289 case 2: 

290 element = _decode_string(reader) 

291 case 3: 

292 element = _decode_document(reader) 

293 case 4: 

294 element = _decode_array(reader) 

295 case 5: 

296 element = _decode_binary(reader) 

297 case 7: 

298 element = _decode_object_id(reader) 

299 case 8: 

300 element = _decode_boolean(reader) 

301 case 16: 

302 element = _decode_int32(reader) 

303 case 18: 

304 element = _decode_int64(reader) 

305 case _: 

306 raise BsonError(f"Bad kind {kind}.") 

307 

308 document.elements.append((name, element)) 

309 

310 return document 

311 

312func _decode_array(reader: _Reader) -> Array: 

313 array = Array() 

314 array.doc = _decode_document(reader) 

315 

316 return array 

317 

318func _decode_binary(reader: _Reader) -> Binary: 

319 length = reader.read_i32() 

320 subtype = reader.read_u8() 

321 data = reader.read_bytes(i64(length)) 

322 

323 return Binary((subtype, data)) 

324 

325func _decode_object_id(reader: _Reader) -> ObjectId: 

326 return ObjectId(reader.read_bytes(12)) 

327 

328func _decode_boolean(reader: _Reader) -> Boolean: 

329 return Boolean(reader.read_u8() == 1) 

330 

331func _decode_int32(reader: _Reader) -> Int32: 

332 return Int32(reader.read_i32()) 

333 

334func _decode_int64(reader: _Reader) -> Int64: 

335 return Int64(reader.read_i64()) 

336 

337func decode(data: bytes) -> Document: 

338 """Decode given BSON document. 

339 

340 """ 

341 

342 return _decode_document(_Reader(data)) 

343 

344func _encode_double(writer: _Writer, name: string, element: Double): 

345 writer.write_u8(1) 

346 writer.write_bytes(name.to_utf8()) 

347 writer.write_u8(0) 

348 writer.write_f64(element.value) 

349 

350func _encode_string(writer: _Writer, name: string, element: String): 

351 writer.write_u8(2) 

352 writer.write_bytes(name.to_utf8()) 

353 writer.write_u8(0) 

354 value = element.value.to_utf8() 

355 writer.write_i32(i32(value.length()) + 1) 

356 writer.write_bytes(value) 

357 writer.write_u8(0) 

358 

359func _encode_document_no_header(writer: _Writer, document: Document): 

360 start = writer.data().length() 

361 writer.write_i32(0) 

362 

363 for name, element in document.elements: 

364 match element: 

365 case Double() as double_element: 

366 _encode_double(writer, name, double_element) 

367 case String() as string_element: 

368 _encode_string(writer, name, string_element) 

369 case Document() as document_element: 

370 _encode_document(writer, name, document_element) 

371 case Array() as array_element: 

372 _encode_array(writer, name, array_element) 

373 case Boolean() as boolean_element: 

374 _encode_boolean(writer, name, boolean_element) 

375 case Int32() as int32_element: 

376 _encode_int32(writer, name, int32_element) 

377 case Int64() as int64_element: 

378 _encode_int64(writer, name, int64_element) 

379 case _: 

380 raise BsonError(f"Bad element '{element}'.") 

381 

382 writer.write_u8(0) 

383 writer.write_i32_at(start, i32(writer.data().length() - start)) 

384 

385func _encode_document(writer: _Writer, name: string, element: Document): 

386 writer.write_u8(3) 

387 writer.write_bytes(name.to_utf8()) 

388 writer.write_u8(0) 

389 _encode_document_no_header(writer, element) 

390 

391func _encode_array(writer: _Writer, name: string, element: Array): 

392 writer.write_u8(4) 

393 writer.write_bytes(name.to_utf8()) 

394 writer.write_u8(0) 

395 _encode_document_no_header(writer, element.doc) 

396 

397func _encode_boolean(writer: _Writer, name: string, element: Boolean): 

398 writer.write_u8(8) 

399 writer.write_bytes(name.to_utf8()) 

400 writer.write_u8(0) 

401 writer.write_u8(u8(1 if element.value else 0)) 

402 

403func _encode_int32(writer: _Writer, name: string, element: Int32): 

404 writer.write_u8(16) 

405 writer.write_bytes(name.to_utf8()) 

406 writer.write_u8(0) 

407 writer.write_i32(element.value) 

408 

409func _encode_int64(writer: _Writer, name: string, element: Int64): 

410 writer.write_u8(18) 

411 writer.write_bytes(name.to_utf8()) 

412 writer.write_u8(0) 

413 writer.write_i64(element.value) 

414 

415func encode(document: Document) -> bytes: 

416 """Encode given BSON document. 

417 

418 """ 

419 

420 writer = _Writer() 

421 _encode_document_no_header(writer, document) 

422 

423 return writer.data() 

424 

425test small_mongodb(): 

426 encoded = ( 

427 b"\x24\x00\x00\x00" 

428 b"\x10" b"whatsmyuri\x00" b"\x01\x00\x00\x00" 

429 b"\x02" b"$db\x00" b"\x06\x00\x00\x00" b"admin\x00" 

430 b"\x00") 

431 

432 document = decode(encoded) 

433 assert document.get("whatsmyuri").int32() == 1 

434 assert document.get("$db").string() == "admin" 

435 

436 encoded_swapped = ( 

437 b"\x24\x00\x00\x00" 

438 b"\x02" b"$db\x00" b"\x06\x00\x00\x00" b"admin\x00" 

439 b"\x10" b"whatsmyuri\x00" b"\x01\x00\x00\x00" 

440 b"\x00") 

441 

442 assert encode(document) in [encoded, encoded_swapped] 

443 

444test mongodb_list_databases(): 

445 encoded = ( 

446 b"\x5a\x00\x00\x00\x01\x6c\x69\x73\x74\x44\x61\x74\x61\x62\x61\x73" 

447 b"\x65\x73\x00\x00\x00\x00\x00\x00\x00\xf0\x3f\x08\x6e\x61\x6d\x65" 

448 b"\x4f\x6e\x6c\x79\x00\x00\x03\x6c\x73\x69\x64\x00\x1e\x00\x00\x00" 

449 b"\x05\x69\x64\x00\x10\x00\x00\x00\x04\xfb\xe9\x2e\x09\xc9\xc8\x45" 

450 b"\x09\xa8\x01\x1d\x0f\x73\xe5\xf7\x46\x00\x02\x24\x64\x62\x00\x06" 

451 b"\x00\x00\x00\x61\x64\x6d\x69\x6e\x00\x00") 

452 

453 document = decode(encoded) 

454 assert document.get("listDatabases").double() == 1.0 

455 assert not document.get("nameOnly").boolean() 

456 assert document.get("lsid").get("id").binary() == ( 

457 4, 

458 b"\xfb\xe9\x2e\x09\xc9\xc8\x45\x09\xa8\x01\x1d\x0f\x73\xe5\xf7\x46") 

459 assert document.get("$db").string() == "admin" 

460 

461test double(): 

462 document = Document([("foo", Double(1.0))]) 

463 encoded = encode(document) 

464 assert encoded == b"\x12\x00\x00\x00\x01foo\x00\x00\x00\x00\x00\x00\x00\xf0?\x00" 

465 decoded = decode(encoded) 

466 assert decoded.get("foo").double() == 1.0 

467 

468test string(): 

469 document = Document([("foo", String("bar"))]) 

470 encoded = encode(document) 

471 assert encoded == b"\x12\x00\x00\x00\x02foo\x00\x04\x00\x00\x00bar\x00\x00" 

472 decoded = decode(encoded) 

473 assert decoded.get("foo").string() == "bar" 

474 

475test document(): 

476 document = Document([("a", Document()), ("b", Boolean(False))]) 

477 encoded = encode(document) 

478 assert encoded == ( 

479 b"\x11\x00\x00\x00\x03a\x00\x05\x00\x00\x00\x00\x08b\x00\x00\x00") 

480 decoded = decode(encoded) 

481 assert decoded.document().length() == 2 

482 assert decoded.get("a").document().length() == 0 

483 assert not decoded.get("b").boolean() 

484 document.elements.append(("foo", String("bar"))) 

485 assert document.get("foo").string() == "bar" 

486 

487 try: 

488 document.get("bar") 

489 assert False 

490 except BsonError as error: 

491 assert error.message == "Element 'bar' not found." 

492 

493test array(): 

494 document = Document([("foo", Array([Int64(5)]))]) 

495 encoded = encode(document) 

496 assert encoded == ( 

497 b"\x1a\x00\x00\x00\x04foo\x00\x10\x00\x00\x00\x120\x00\x05\x00" 

498 b"\x00\x00\x00\x00\x00\x00\x00\x00") 

499 decoded = decode(encoded) 

500 assert decoded.get("foo").array().length() == 1 

501 assert decoded.get("foo").at(0).int64() == 5