Coverage for src/lib.mys : 94%

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
1c"""source-before-namespace
2#include <cstring>
3"""
5class BsonError(Error):
6 message: string
8class _Reader:
9 _data: bytes
10 _pos: i64
12 func __init__(self, data: bytes):
13 self._data = data
14 self._pos = 0
16 func available(self) -> bool:
17 return self._pos < self._data.length()
19 func read_u8(self) -> u8:
20 if not self.available():
21 raise BsonError("No more data available.")
23 value = self._data[self._pos]
24 self._pos += 1
26 return value
28 func read_u16(self) -> u16:
29 return u16(self.read_u8()) | (u16(self.read_u8()) << 8)
31 func read_u32(self) -> u32:
32 return u32(self.read_u16()) | (u32(self.read_u16()) << 16)
34 func read_u64(self) -> u64:
35 return u64(self.read_u32()) | (u64(self.read_u32()) << 32)
37 func read_i32(self) -> i32:
38 return i32(self.read_u32())
40 func read_i64(self) -> i64:
41 return i64(self.read_u64())
43 func read_f64(self) -> f64:
44 value = 0.0
45 ivalue = self.read_u64()
46 c"memcpy(&value, &ivalue, sizeof(value));"
48 return value
50 func read_bytes(self, count: i64) -> bytes:
51 data = b""
53 for _ in range(count):
54 data += self.read_u8()
56 return data
58class _Writer:
59 _data: bytes
61 func __init__(self):
62 self._data = b""
64 func data(self) -> bytes:
65 return self._data
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)
73 func write_u8(self, value: u8):
74 self._data += value
76 func write_u16(self, value: u16):
77 self.write_u8(u8(value))
78 self.write_u8(u8(value >> 8))
80 func write_u32(self, value: u32):
81 self.write_u16(u16(value))
82 self.write_u16(u16(value >> 16))
84 func write_u64(self, value: u64):
85 self.write_u32(u32(value))
86 self.write_u32(u32(value >> 32))
88 func write_i32(self, value: i32):
89 self.write_u32(u32(value))
91 func write_i64(self, value: i64):
92 self.write_u64(u64(value))
94 func write_f64(self, value: f64):
95 ivalue = 0
96 c"memcpy(&ivalue, &value, sizeof(ivalue));"
97 self.write_i64(ivalue)
99 func write_bytes(self, value: bytes):
100 self._data += value
102trait Element:
103 """An element in a BSON document.
105 """
107 func double(self) -> f64:
108 raise BsonError("Not a double.")
110 func string(self) -> string:
111 raise BsonError("Not a string.")
113 func document(self) -> [(string, Element)]:
114 raise BsonError("Not a document.")
116 func get(self, name: string) -> Element:
117 raise BsonError("Not a document.")
119 func array(self) -> [Element]:
120 raise BsonError("Not an array.")
122 func at(self, index: i64) -> Element:
123 raise BsonError("Not an array.")
125 func binary(self) -> (u8, bytes):
126 raise BsonError("Not a binary.")
128 func boolean(self) -> bool:
129 raise BsonError("Not a boolean.")
131 func int32(self) -> i32:
132 raise BsonError("Not an int32.")
134 func int64(self) -> i64:
135 raise BsonError("Not an int64.")
137class Double(Element):
138 """A double.
140 """
142 value: f64
144 func double(self) -> f64:
145 return self.value
147class String(Element):
148 """A string.
150 """
152 value: string
154 func string(self) -> string:
155 return self.value
157class Document(Element):
158 """A document.
160 """
162 elements: [(string, Element)]
164 func __init__(self, elements: [(string, Element)] = []):
165 self.elements = elements
167 func document(self) -> [(string, Element)]:
168 return self.elements
170 func get(self, name: string) -> Element:
171 for key, element in self.elements:
172 if key == name:
173 return element
175 raise BsonError(f"Element '{name}' not found.")
177class Array(Element):
178 """An array.
180 """
182 doc: Document
184 func __init__(self, elements: [Element] = []):
185 self.doc = Document()
187 for i, element in enumerate(elements):
188 self.doc.elements.append((str(i), element))
190 func append(self, element: Element):
191 index = str(self.doc.elements.length())
192 self.doc.elements.append((index, element))
194 func array(self) -> [Element]:
195 return [element for _, element in self.doc.elements]
197 func at(self, index: i64) -> Element:
198 return self.doc.elements[index][1]
200class Binary(Element):
201 """A binary blob.
203 """
205 value: (u8, bytes)
207 func binary(self) -> (u8, bytes):
208 return self.value
210class ObjectId(Element):
211 """An object id.
213 """
215 value: bytes
217 func object_id(self) -> bytes:
218 return self.value
220class Boolean(Element):
221 """A boolean.
223 """
225 value: bool
227 func boolean(self) -> bool:
228 return self.value
230class Int32(Element):
231 """An i32.
233 """
235 value: i32
237 func int32(self) -> i32:
238 return self.value
240class Int64(Element):
241 """An i64.
243 """
245 value: i64
247 func int64(self) -> i64:
248 return self.value
250func _decode_cstring(reader: _Reader) -> string:
251 data = b""
253 while True:
254 value = reader.read_u8()
256 if value == 0:
257 break
259 data += value
261 return string(data)
263func _decode_double(reader: _Reader) -> Double:
264 return Double(reader.read_f64())
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()
271 return element
273func _decode_document(reader: _Reader) -> Document:
274 reader.read_i32()
275 document = Document()
277 while True:
278 kind = reader.read_u8()
280 if kind == 0:
281 break
283 name = _decode_cstring(reader)
284 element: Element? = None
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}.")
308 document.elements.append((name, element))
310 return document
312func _decode_array(reader: _Reader) -> Array:
313 array = Array()
314 array.doc = _decode_document(reader)
316 return array
318func _decode_binary(reader: _Reader) -> Binary:
319 length = reader.read_i32()
320 subtype = reader.read_u8()
321 data = reader.read_bytes(i64(length))
323 return Binary((subtype, data))
325func _decode_object_id(reader: _Reader) -> ObjectId:
326 return ObjectId(reader.read_bytes(12))
328func _decode_boolean(reader: _Reader) -> Boolean:
329 return Boolean(reader.read_u8() == 1)
331func _decode_int32(reader: _Reader) -> Int32:
332 return Int32(reader.read_i32())
334func _decode_int64(reader: _Reader) -> Int64:
335 return Int64(reader.read_i64())
337func decode(data: bytes) -> Document:
338 """Decode given BSON document.
340 """
342 return _decode_document(_Reader(data))
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)
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)
359func _encode_document_no_header(writer: _Writer, document: Document):
360 start = writer.data().length()
361 writer.write_i32(0)
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}'.")
382 writer.write_u8(0)
383 writer.write_i32_at(start, i32(writer.data().length() - start))
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)
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)
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))
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)
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)
415func encode(document: Document) -> bytes:
416 """Encode given BSON document.
418 """
420 writer = _Writer()
421 _encode_document_no_header(writer, document)
423 return writer.data()
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")
432 document = decode(encoded)
433 assert document.get("whatsmyuri").int32() == 1
434 assert document.get("$db").string() == "admin"
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")
442 assert encode(document) in [encoded, encoded_swapped]
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")
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"
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
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"
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"
487 try:
488 document.get("bar")
489 assert False
490 except BsonError as error:
491 assert error.message == "Element 'bar' not found."
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