zlmdb 25.12.3__cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (78) hide show
  1. zlmdb/__init__.py +416 -0
  2. zlmdb/_database.py +990 -0
  3. zlmdb/_errors.py +31 -0
  4. zlmdb/_flatc/__init__.py +112 -0
  5. zlmdb/_flatc/bin/flatc +0 -0
  6. zlmdb/_lmdb_vendor/__init__.py +37 -0
  7. zlmdb/_lmdb_vendor/__main__.py +25 -0
  8. zlmdb/_lmdb_vendor/_config.py +10 -0
  9. zlmdb/_lmdb_vendor/_lmdb_cffi.cpython-312-aarch64-linux-gnu.so +0 -0
  10. zlmdb/_lmdb_vendor/cffi.py +2606 -0
  11. zlmdb/_lmdb_vendor/tool.py +670 -0
  12. zlmdb/_meta.py +27 -0
  13. zlmdb/_pmap.py +1667 -0
  14. zlmdb/_schema.py +137 -0
  15. zlmdb/_transaction.py +181 -0
  16. zlmdb/_types.py +1596 -0
  17. zlmdb/_version.py +27 -0
  18. zlmdb/cli.py +41 -0
  19. zlmdb/flatbuffers/__init__.py +60 -0
  20. zlmdb/flatbuffers/_git_version.py +24 -0
  21. zlmdb/flatbuffers/_version.py +17 -0
  22. zlmdb/flatbuffers/builder.py +824 -0
  23. zlmdb/flatbuffers/compat.py +89 -0
  24. zlmdb/flatbuffers/encode.py +43 -0
  25. zlmdb/flatbuffers/flexbuffers.py +1570 -0
  26. zlmdb/flatbuffers/number_types.py +182 -0
  27. zlmdb/flatbuffers/packer.py +42 -0
  28. zlmdb/flatbuffers/reflection/AdvancedFeatures.py +10 -0
  29. zlmdb/flatbuffers/reflection/BaseType.py +25 -0
  30. zlmdb/flatbuffers/reflection/Enum.py +252 -0
  31. zlmdb/flatbuffers/reflection/EnumVal.py +144 -0
  32. zlmdb/flatbuffers/reflection/Field.py +325 -0
  33. zlmdb/flatbuffers/reflection/KeyValue.py +84 -0
  34. zlmdb/flatbuffers/reflection/Object.py +260 -0
  35. zlmdb/flatbuffers/reflection/RPCCall.py +195 -0
  36. zlmdb/flatbuffers/reflection/Schema.py +301 -0
  37. zlmdb/flatbuffers/reflection/SchemaFile.py +112 -0
  38. zlmdb/flatbuffers/reflection/Service.py +213 -0
  39. zlmdb/flatbuffers/reflection/Type.py +148 -0
  40. zlmdb/flatbuffers/reflection/__init__.py +0 -0
  41. zlmdb/flatbuffers/reflection.bfbs +0 -0
  42. zlmdb/flatbuffers/reflection.fbs +156 -0
  43. zlmdb/flatbuffers/table.py +129 -0
  44. zlmdb/flatbuffers/util.py +47 -0
  45. zlmdb/py.typed +0 -0
  46. zlmdb/tests/conftest.py +9 -0
  47. zlmdb/tests/lmdb/__init__.py +0 -0
  48. zlmdb/tests/lmdb/address_book.py +287 -0
  49. zlmdb/tests/lmdb/crash_test.py +339 -0
  50. zlmdb/tests/lmdb/cursor_test.py +333 -0
  51. zlmdb/tests/lmdb/env_test.py +919 -0
  52. zlmdb/tests/lmdb/getmulti_test.py +92 -0
  53. zlmdb/tests/lmdb/iteration_test.py +258 -0
  54. zlmdb/tests/lmdb/package_test.py +70 -0
  55. zlmdb/tests/lmdb/test_lmdb.py +188 -0
  56. zlmdb/tests/lmdb/testlib.py +185 -0
  57. zlmdb/tests/lmdb/tool_test.py +60 -0
  58. zlmdb/tests/lmdb/txn_test.py +575 -0
  59. zlmdb/tests/orm/MNodeLog.py +853 -0
  60. zlmdb/tests/orm/__init__.py +0 -0
  61. zlmdb/tests/orm/_schema_fbs.py +215 -0
  62. zlmdb/tests/orm/_schema_mnode_log.py +1202 -0
  63. zlmdb/tests/orm/_schema_py2.py +250 -0
  64. zlmdb/tests/orm/_schema_py3.py +307 -0
  65. zlmdb/tests/orm/_test_flatbuffers.py +144 -0
  66. zlmdb/tests/orm/_test_serialization.py +144 -0
  67. zlmdb/tests/orm/test_basic.py +217 -0
  68. zlmdb/tests/orm/test_etcd.py +275 -0
  69. zlmdb/tests/orm/test_pmap_indexes.py +466 -0
  70. zlmdb/tests/orm/test_pmap_types.py +90 -0
  71. zlmdb/tests/orm/test_pmaps.py +295 -0
  72. zlmdb/tests/orm/test_select.py +619 -0
  73. zlmdb-25.12.3.dist-info/METADATA +277 -0
  74. zlmdb-25.12.3.dist-info/RECORD +78 -0
  75. zlmdb-25.12.3.dist-info/WHEEL +6 -0
  76. zlmdb-25.12.3.dist-info/entry_points.txt +3 -0
  77. zlmdb-25.12.3.dist-info/licenses/LICENSE +167 -0
  78. zlmdb-25.12.3.dist-info/licenses/NOTICE +41 -0
@@ -0,0 +1,1570 @@
1
+ # Lint as: python3
2
+ # Copyright 2020 Google Inc. All rights reserved.
3
+ #
4
+ # Licensed under the Apache License, Version 2.0 (the "License");
5
+ # you may not use this file except in compliance with the License.
6
+ # You may obtain a copy of the License at
7
+ #
8
+ # http://www.apache.org/licenses/LICENSE-2.0
9
+ #
10
+ # Unless required by applicable law or agreed to in writing, software
11
+ # distributed under the License is distributed on an "AS IS" BASIS,
12
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13
+ # See the License for the specific language governing permissions and
14
+ # limitations under the License.
15
+ """Implementation of FlexBuffers binary format.
16
+
17
+ For more info check https://google.github.io/flatbuffers/flexbuffers.html and
18
+ corresponding C++ implementation at
19
+ https://github.com/google/flatbuffers/blob/master/include/flatbuffers/flexbuffers.h
20
+ """
21
+
22
+ # pylint: disable=invalid-name
23
+ # TODO(dkovalev): Add type hints everywhere, so tools like pytypes could work.
24
+
25
+ import array
26
+ import contextlib
27
+ import enum
28
+ import struct
29
+
30
+ __all__ = ("Type", "Builder", "GetRoot", "Dumps", "Loads")
31
+
32
+
33
+ class BitWidth(enum.IntEnum):
34
+ """Supported bit widths of value types.
35
+
36
+ These are used in the lower 2 bits of a type field to determine the size of
37
+ the elements (and or size field) of the item pointed to (e.g. vector).
38
+ """
39
+
40
+ W8 = 0 # 2^0 = 1 byte
41
+ W16 = 1 # 2^1 = 2 bytes
42
+ W32 = 2 # 2^2 = 4 bytes
43
+ W64 = 3 # 2^3 = 8 bytes
44
+
45
+ @staticmethod
46
+ def U(value):
47
+ """Returns the minimum `BitWidth` to encode unsigned integer value."""
48
+ assert value >= 0
49
+
50
+ if value < (1 << 8):
51
+ return BitWidth.W8
52
+ elif value < (1 << 16):
53
+ return BitWidth.W16
54
+ elif value < (1 << 32):
55
+ return BitWidth.W32
56
+ elif value < (1 << 64):
57
+ return BitWidth.W64
58
+ else:
59
+ raise ValueError("value is too big to encode: %s" % value)
60
+
61
+ @staticmethod
62
+ def I(value):
63
+ """Returns the minimum `BitWidth` to encode signed integer value."""
64
+ # -2^(n-1) <= value < 2^(n-1)
65
+ # -2^n <= 2 * value < 2^n
66
+ # 2 * value < 2^n, when value >= 0 or 2 * (-value) <= 2^n, when value < 0
67
+ # 2 * value < 2^n, when value >= 0 or 2 * (-value) - 1 < 2^n, when value < 0
68
+ #
69
+ # if value >= 0:
70
+ # return BitWidth.U(2 * value)
71
+ # else:
72
+ # return BitWidth.U(2 * (-value) - 1) # ~x = -x - 1
73
+ value *= 2
74
+ return BitWidth.U(value if value >= 0 else ~value)
75
+
76
+ @staticmethod
77
+ def F(value):
78
+ """Returns the `BitWidth` to encode floating point value."""
79
+ if struct.unpack("f", struct.pack("f", value))[0] == value:
80
+ return BitWidth.W32
81
+ return BitWidth.W64
82
+
83
+ @staticmethod
84
+ def B(byte_width):
85
+ return {1: BitWidth.W8, 2: BitWidth.W16, 4: BitWidth.W32, 8: BitWidth.W64}[
86
+ byte_width
87
+ ]
88
+
89
+
90
+ I = {1: "b", 2: "h", 4: "i", 8: "q"} # Integer formats
91
+ U = {1: "B", 2: "H", 4: "I", 8: "Q"} # Unsigned integer formats
92
+ F = {4: "f", 8: "d"} # Floating point formats
93
+
94
+
95
+ def _Unpack(fmt, buf):
96
+ return struct.unpack(fmt[len(buf)], buf)[0]
97
+
98
+
99
+ def _UnpackVector(fmt, buf, length):
100
+ byte_width = len(buf) // length
101
+ return struct.unpack("%d%s" % (length, fmt[byte_width]), buf)
102
+
103
+
104
+ def _Pack(fmt, value, byte_width):
105
+ return struct.pack(fmt[byte_width], value)
106
+
107
+
108
+ def _PackVector(fmt, values, byte_width):
109
+ return struct.pack("%d%s" % (len(values), fmt[byte_width]), *values)
110
+
111
+
112
+ def _Mutate(fmt, buf, value, byte_width, value_bit_width):
113
+ if (1 << value_bit_width) <= byte_width:
114
+ buf[:byte_width] = _Pack(fmt, value, byte_width)
115
+ return True
116
+ return False
117
+
118
+
119
+ # Computes how many bytes you'd have to pad to be able to write an
120
+ # "scalar_size" scalar if the buffer had grown to "buf_size",
121
+ # "scalar_size" is a power of two.
122
+ def _PaddingBytes(buf_size, scalar_size):
123
+ # ((buf_size + (scalar_size - 1)) // scalar_size) * scalar_size - buf_size
124
+ return -buf_size & (scalar_size - 1)
125
+
126
+
127
+ def _ShiftSlice(s, offset, length):
128
+ start = offset + (0 if s.start is None else s.start)
129
+ stop = offset + (length if s.stop is None else s.stop)
130
+ return slice(start, stop, s.step)
131
+
132
+
133
+ # https://en.cppreference.com/w/cpp/algorithm/lower_bound
134
+ def _LowerBound(values, value, pred):
135
+ """Implementation of C++ std::lower_bound() algorithm."""
136
+ first, last = 0, len(values)
137
+ count = last - first
138
+ while count > 0:
139
+ i = first
140
+ step = count // 2
141
+ i += step
142
+ if pred(values[i], value):
143
+ i += 1
144
+ first = i
145
+ count -= step + 1
146
+ else:
147
+ count = step
148
+ return first
149
+
150
+
151
+ # https://en.cppreference.com/w/cpp/algorithm/binary_search
152
+ def _BinarySearch(values, value, pred=lambda x, y: x < y):
153
+ """Implementation of C++ std::binary_search() algorithm."""
154
+ index = _LowerBound(values, value, pred)
155
+ if index != len(values) and not pred(value, values[index]):
156
+ return index
157
+ return -1
158
+
159
+
160
+ class Type(enum.IntEnum):
161
+ """Supported types of encoded data.
162
+
163
+ These are used as the upper 6 bits of a type field to indicate the actual
164
+ type.
165
+ """
166
+
167
+ NULL = 0
168
+ INT = 1
169
+ UINT = 2
170
+ FLOAT = 3
171
+ # Types above stored inline, types below store an offset.
172
+ KEY = 4
173
+ STRING = 5
174
+ INDIRECT_INT = 6
175
+ INDIRECT_UINT = 7
176
+ INDIRECT_FLOAT = 8
177
+ MAP = 9
178
+ VECTOR = 10 # Untyped.
179
+
180
+ VECTOR_INT = 11 # Typed any size (stores no type table).
181
+ VECTOR_UINT = 12
182
+ VECTOR_FLOAT = 13
183
+ VECTOR_KEY = 14
184
+ # DEPRECATED, use VECTOR or VECTOR_KEY instead.
185
+ # Read test.cpp/FlexBuffersDeprecatedTest() for details on why.
186
+ VECTOR_STRING_DEPRECATED = 15
187
+
188
+ VECTOR_INT2 = 16 # Typed tuple (no type table, no size field).
189
+ VECTOR_UINT2 = 17
190
+ VECTOR_FLOAT2 = 18
191
+ VECTOR_INT3 = 19 # Typed triple (no type table, no size field).
192
+ VECTOR_UINT3 = 20
193
+ VECTOR_FLOAT3 = 21
194
+ VECTOR_INT4 = 22 # Typed quad (no type table, no size field).
195
+ VECTOR_UINT4 = 23
196
+ VECTOR_FLOAT4 = 24
197
+
198
+ BLOB = 25
199
+ BOOL = 26
200
+ VECTOR_BOOL = 36 # To do the same type of conversion of type to vector type
201
+
202
+ @staticmethod
203
+ def Pack(type_, bit_width):
204
+ return (int(type_) << 2) | bit_width
205
+
206
+ @staticmethod
207
+ def Unpack(packed_type):
208
+ return 1 << (packed_type & 0b11), Type(packed_type >> 2)
209
+
210
+ @staticmethod
211
+ def IsInline(type_):
212
+ return type_ <= Type.FLOAT or type_ == Type.BOOL
213
+
214
+ @staticmethod
215
+ def IsTypedVector(type_):
216
+ return (
217
+ Type.VECTOR_INT <= type_ <= Type.VECTOR_STRING_DEPRECATED
218
+ or type_ == Type.VECTOR_BOOL
219
+ )
220
+
221
+ @staticmethod
222
+ def IsTypedVectorElementType(type_):
223
+ return Type.INT <= type_ <= Type.STRING or type_ == Type.BOOL
224
+
225
+ @staticmethod
226
+ def ToTypedVectorElementType(type_):
227
+ if not Type.IsTypedVector(type_):
228
+ raise ValueError("must be typed vector type")
229
+
230
+ return Type(type_ - Type.VECTOR_INT + Type.INT)
231
+
232
+ @staticmethod
233
+ def IsFixedTypedVector(type_):
234
+ return Type.VECTOR_INT2 <= type_ <= Type.VECTOR_FLOAT4
235
+
236
+ @staticmethod
237
+ def IsFixedTypedVectorElementType(type_):
238
+ return Type.INT <= type_ <= Type.FLOAT
239
+
240
+ @staticmethod
241
+ def ToFixedTypedVectorElementType(type_):
242
+ if not Type.IsFixedTypedVector(type_):
243
+ raise ValueError("must be fixed typed vector type")
244
+
245
+ # 3 types each, starting from length 2.
246
+ fixed_type = type_ - Type.VECTOR_INT2
247
+ return Type(fixed_type % 3 + Type.INT), fixed_type // 3 + 2
248
+
249
+ @staticmethod
250
+ def ToTypedVector(element_type, fixed_len=0):
251
+ """Converts element type to corresponding vector type.
252
+
253
+ Args:
254
+ element_type: vector element type
255
+ fixed_len: number of elements: 0 for typed vector; 2, 3, or 4 for fixed
256
+ typed vector.
257
+
258
+ Returns:
259
+ Typed vector type or fixed typed vector type.
260
+ """
261
+ if fixed_len == 0:
262
+ if not Type.IsTypedVectorElementType(element_type):
263
+ raise ValueError("must be typed vector element type")
264
+ else:
265
+ if not Type.IsFixedTypedVectorElementType(element_type):
266
+ raise ValueError("must be fixed typed vector element type")
267
+
268
+ offset = element_type - Type.INT
269
+ if fixed_len == 0:
270
+ return Type(offset + Type.VECTOR_INT) # TypedVector
271
+ elif fixed_len == 2:
272
+ return Type(offset + Type.VECTOR_INT2) # FixedTypedVector
273
+ elif fixed_len == 3:
274
+ return Type(offset + Type.VECTOR_INT3) # FixedTypedVector
275
+ elif fixed_len == 4:
276
+ return Type(offset + Type.VECTOR_INT4) # FixedTypedVector
277
+ else:
278
+ raise ValueError("unsupported fixed_len: %s" % fixed_len)
279
+
280
+
281
+ class Buf:
282
+ """Class to access underlying buffer object starting from the given offset."""
283
+
284
+ def __init__(self, buf, offset):
285
+ self._buf = buf
286
+ self._offset = offset if offset >= 0 else len(buf) + offset
287
+ self._length = len(buf) - self._offset
288
+
289
+ def __getitem__(self, key):
290
+ if isinstance(key, slice):
291
+ return self._buf[_ShiftSlice(key, self._offset, self._length)]
292
+ elif isinstance(key, int):
293
+ return self._buf[self._offset + key]
294
+ else:
295
+ raise TypeError("invalid key type")
296
+
297
+ def __setitem__(self, key, value):
298
+ if isinstance(key, slice):
299
+ self._buf[_ShiftSlice(key, self._offset, self._length)] = value
300
+ elif isinstance(key, int):
301
+ self._buf[self._offset + key] = key
302
+ else:
303
+ raise TypeError("invalid key type")
304
+
305
+ def __repr__(self):
306
+ return "buf[%d:]" % self._offset
307
+
308
+ def Find(self, sub):
309
+ """Returns the lowest index where the sub subsequence is found."""
310
+ return self._buf[self._offset :].find(sub)
311
+
312
+ def Slice(self, offset):
313
+ """Returns new `Buf` which starts from the given offset."""
314
+ return Buf(self._buf, self._offset + offset)
315
+
316
+ def Indirect(self, offset, byte_width):
317
+ """Return new `Buf` based on the encoded offset (indirect encoding)."""
318
+ return self.Slice(offset - _Unpack(U, self[offset : offset + byte_width]))
319
+
320
+
321
+ class Object:
322
+ """Base class for all non-trivial data accessors."""
323
+
324
+ __slots__ = "_buf", "_byte_width"
325
+
326
+ def __init__(self, buf, byte_width):
327
+ self._buf = buf
328
+ self._byte_width = byte_width
329
+
330
+ @property
331
+ def ByteWidth(self):
332
+ return self._byte_width
333
+
334
+
335
+ class Sized(Object):
336
+ """Base class for all data accessors which need to read encoded size."""
337
+
338
+ __slots__ = ("_size",)
339
+
340
+ def __init__(self, buf, byte_width, size=0):
341
+ super().__init__(buf, byte_width)
342
+ if size == 0:
343
+ self._size = _Unpack(U, self.SizeBytes)
344
+ else:
345
+ self._size = size
346
+
347
+ @property
348
+ def SizeBytes(self):
349
+ return self._buf[-self._byte_width : 0]
350
+
351
+ def __len__(self):
352
+ return self._size
353
+
354
+
355
+ class Blob(Sized):
356
+ """Data accessor for the encoded blob bytes."""
357
+
358
+ __slots__ = ()
359
+
360
+ @property
361
+ def Bytes(self):
362
+ return self._buf[0 : len(self)]
363
+
364
+ def __repr__(self):
365
+ return "Blob(%s, size=%d)" % (self._buf, len(self))
366
+
367
+
368
+ class String(Sized):
369
+ """Data accessor for the encoded string bytes."""
370
+
371
+ __slots__ = ()
372
+
373
+ @property
374
+ def Bytes(self):
375
+ return self._buf[0 : len(self)]
376
+
377
+ def Mutate(self, value):
378
+ """Mutates underlying string bytes in place.
379
+
380
+ Args:
381
+ value: New string to replace the existing one. New string must have less
382
+ or equal UTF-8-encoded bytes than the existing one to successfully
383
+ mutate underlying byte buffer.
384
+
385
+ Returns:
386
+ Whether the value was mutated or not.
387
+ """
388
+ encoded = value.encode("utf-8")
389
+ n = len(encoded)
390
+ if n <= len(self):
391
+ self._buf[-self._byte_width : 0] = _Pack(U, n, self._byte_width)
392
+ self._buf[0:n] = encoded
393
+ self._buf[n : len(self)] = bytearray(len(self) - n)
394
+ return True
395
+ return False
396
+
397
+ def __str__(self):
398
+ return self.Bytes.decode("utf-8")
399
+
400
+ def __repr__(self):
401
+ return "String(%s, size=%d)" % (self._buf, len(self))
402
+
403
+
404
+ class Key(Object):
405
+ """Data accessor for the encoded key bytes."""
406
+
407
+ __slots__ = ()
408
+
409
+ def __init__(self, buf, byte_width):
410
+ assert byte_width == 1
411
+ super().__init__(buf, byte_width)
412
+
413
+ @property
414
+ def Bytes(self):
415
+ return self._buf[0 : len(self)]
416
+
417
+ def __len__(self):
418
+ return self._buf.Find(0)
419
+
420
+ def __str__(self):
421
+ return self.Bytes.decode("ascii")
422
+
423
+ def __repr__(self):
424
+ return "Key(%s, size=%d)" % (self._buf, len(self))
425
+
426
+
427
+ class Vector(Sized):
428
+ """Data accessor for the encoded vector bytes."""
429
+
430
+ __slots__ = ()
431
+
432
+ def __getitem__(self, index):
433
+ if index < 0 or index >= len(self):
434
+ raise IndexError(
435
+ "vector index %s is out of [0, %d) range" % (index, len(self))
436
+ )
437
+
438
+ packed_type = self._buf[len(self) * self._byte_width + index]
439
+ buf = self._buf.Slice(index * self._byte_width)
440
+ return Ref.PackedType(buf, self._byte_width, packed_type)
441
+
442
+ @property
443
+ def Value(self):
444
+ """Returns the underlying encoded data as a list object."""
445
+ return [e.Value for e in self]
446
+
447
+ def __repr__(self):
448
+ return "Vector(%s, byte_width=%d, size=%d)" % (
449
+ self._buf,
450
+ self._byte_width,
451
+ self._size,
452
+ )
453
+
454
+
455
+ class TypedVector(Sized):
456
+ """Data accessor for the encoded typed vector or fixed typed vector bytes."""
457
+
458
+ __slots__ = "_element_type", "_size"
459
+
460
+ def __init__(self, buf, byte_width, element_type, size=0):
461
+ super().__init__(buf, byte_width, size)
462
+
463
+ if element_type == Type.STRING:
464
+ # These can't be accessed as strings, since we don't know the bit-width
465
+ # of the size field, see the declaration of
466
+ # FBT_VECTOR_STRING_DEPRECATED above for details.
467
+ # We change the type here to be keys, which are a subtype of strings,
468
+ # and will ignore the size field. This will truncate strings with
469
+ # embedded nulls.
470
+ element_type = Type.KEY
471
+
472
+ self._element_type = element_type
473
+
474
+ @property
475
+ def Bytes(self):
476
+ return self._buf[: self._byte_width * len(self)]
477
+
478
+ @property
479
+ def ElementType(self):
480
+ return self._element_type
481
+
482
+ def __getitem__(self, index):
483
+ if index < 0 or index >= len(self):
484
+ raise IndexError(
485
+ "vector index %s is out of [0, %d) range" % (index, len(self))
486
+ )
487
+
488
+ buf = self._buf.Slice(index * self._byte_width)
489
+ return Ref(buf, self._byte_width, 1, self._element_type)
490
+
491
+ @property
492
+ def Value(self):
493
+ """Returns underlying data as list object."""
494
+ if not self:
495
+ return []
496
+
497
+ if self._element_type is Type.BOOL:
498
+ return [bool(e) for e in _UnpackVector(U, self.Bytes, len(self))]
499
+ elif self._element_type is Type.INT:
500
+ return list(_UnpackVector(I, self.Bytes, len(self)))
501
+ elif self._element_type is Type.UINT:
502
+ return list(_UnpackVector(U, self.Bytes, len(self)))
503
+ elif self._element_type is Type.FLOAT:
504
+ return list(_UnpackVector(F, self.Bytes, len(self)))
505
+ elif self._element_type is Type.KEY:
506
+ return [e.AsKey for e in self]
507
+ elif self._element_type is Type.STRING:
508
+ return [e.AsString for e in self]
509
+ else:
510
+ raise TypeError("unsupported element_type: %s" % self._element_type)
511
+
512
+ def __repr__(self):
513
+ return "TypedVector(%s, byte_width=%d, element_type=%s, size=%d)" % (
514
+ self._buf,
515
+ self._byte_width,
516
+ self._element_type,
517
+ self._size,
518
+ )
519
+
520
+
521
+ class Map(Vector):
522
+ """Data accessor for the encoded map bytes."""
523
+
524
+ @staticmethod
525
+ def CompareKeys(a, b):
526
+ if isinstance(a, Ref):
527
+ a = a.AsKeyBytes
528
+ if isinstance(b, Ref):
529
+ b = b.AsKeyBytes
530
+ return a < b
531
+
532
+ def __getitem__(self, key):
533
+ if isinstance(key, int):
534
+ return super().__getitem__(key)
535
+
536
+ index = _BinarySearch(self.Keys, key.encode("ascii"), self.CompareKeys)
537
+ if index != -1:
538
+ return super().__getitem__(index)
539
+
540
+ raise KeyError(key)
541
+
542
+ @property
543
+ def Keys(self):
544
+ byte_width = _Unpack(U, self._buf[-2 * self._byte_width : -self._byte_width])
545
+ buf = self._buf.Indirect(-3 * self._byte_width, self._byte_width)
546
+ return TypedVector(buf, byte_width, Type.KEY)
547
+
548
+ @property
549
+ def Values(self):
550
+ return Vector(self._buf, self._byte_width)
551
+
552
+ @property
553
+ def Value(self):
554
+ return {k.Value: v.Value for k, v in zip(self.Keys, self.Values)}
555
+
556
+ def __repr__(self):
557
+ return "Map(%s, size=%d)" % (self._buf, len(self))
558
+
559
+
560
+ class Ref:
561
+ """Data accessor for the encoded data bytes."""
562
+
563
+ __slots__ = "_buf", "_parent_width", "_byte_width", "_type"
564
+
565
+ @staticmethod
566
+ def PackedType(buf, parent_width, packed_type):
567
+ byte_width, type_ = Type.Unpack(packed_type)
568
+ return Ref(buf, parent_width, byte_width, type_)
569
+
570
+ def __init__(self, buf, parent_width, byte_width, type_):
571
+ self._buf = buf
572
+ self._parent_width = parent_width
573
+ self._byte_width = byte_width
574
+ self._type = type_
575
+
576
+ def __repr__(self):
577
+ return "Ref(%s, parent_width=%d, byte_width=%d, type_=%s)" % (
578
+ self._buf,
579
+ self._parent_width,
580
+ self._byte_width,
581
+ self._type,
582
+ )
583
+
584
+ @property
585
+ def _Bytes(self):
586
+ return self._buf[: self._parent_width]
587
+
588
+ def _ConvertError(self, target_type):
589
+ raise TypeError("cannot convert %s to %s" % (self._type, target_type))
590
+
591
+ def _Indirect(self):
592
+ return self._buf.Indirect(0, self._parent_width)
593
+
594
+ @property
595
+ def IsNull(self):
596
+ return self._type is Type.NULL
597
+
598
+ @property
599
+ def IsBool(self):
600
+ return self._type is Type.BOOL
601
+
602
+ @property
603
+ def AsBool(self):
604
+ if self._type is Type.BOOL:
605
+ return bool(_Unpack(U, self._Bytes))
606
+ else:
607
+ return self.AsInt != 0
608
+
609
+ def MutateBool(self, value):
610
+ """Mutates underlying boolean value bytes in place.
611
+
612
+ Args:
613
+ value: New boolean value.
614
+
615
+ Returns:
616
+ Whether the value was mutated or not.
617
+ """
618
+ return self.IsBool and _Mutate(
619
+ U, self._buf, value, self._parent_width, BitWidth.W8
620
+ )
621
+
622
+ @property
623
+ def IsNumeric(self):
624
+ return self.IsInt or self.IsFloat
625
+
626
+ @property
627
+ def IsInt(self):
628
+ return self._type in (
629
+ Type.INT,
630
+ Type.INDIRECT_INT,
631
+ Type.UINT,
632
+ Type.INDIRECT_UINT,
633
+ )
634
+
635
+ @property
636
+ def AsInt(self):
637
+ """Returns current reference as integer value."""
638
+ if self.IsNull:
639
+ return 0
640
+ elif self.IsBool:
641
+ return int(self.AsBool)
642
+ elif self._type is Type.INT:
643
+ return _Unpack(I, self._Bytes)
644
+ elif self._type is Type.INDIRECT_INT:
645
+ return _Unpack(I, self._Indirect()[: self._byte_width])
646
+ if self._type is Type.UINT:
647
+ return _Unpack(U, self._Bytes)
648
+ elif self._type is Type.INDIRECT_UINT:
649
+ return _Unpack(U, self._Indirect()[: self._byte_width])
650
+ elif self.IsString:
651
+ return len(self.AsString)
652
+ elif self.IsKey:
653
+ return len(self.AsKey)
654
+ elif self.IsBlob:
655
+ return len(self.AsBlob)
656
+ elif self.IsVector:
657
+ return len(self.AsVector)
658
+ elif self.IsTypedVector:
659
+ return len(self.AsTypedVector)
660
+ elif self.IsFixedTypedVector:
661
+ return len(self.AsFixedTypedVector)
662
+ else:
663
+ raise self._ConvertError(Type.INT)
664
+
665
+ def MutateInt(self, value):
666
+ """Mutates underlying integer value bytes in place.
667
+
668
+ Args:
669
+ value: New integer value. It must fit to the byte size of the existing
670
+ encoded value.
671
+
672
+ Returns:
673
+ Whether the value was mutated or not.
674
+ """
675
+ if self._type is Type.INT:
676
+ return _Mutate(I, self._buf, value, self._parent_width, BitWidth.I(value))
677
+ elif self._type is Type.INDIRECT_INT:
678
+ return _Mutate(
679
+ I, self._Indirect(), value, self._byte_width, BitWidth.I(value)
680
+ )
681
+ elif self._type is Type.UINT:
682
+ return _Mutate(U, self._buf, value, self._parent_width, BitWidth.U(value))
683
+ elif self._type is Type.INDIRECT_UINT:
684
+ return _Mutate(
685
+ U, self._Indirect(), value, self._byte_width, BitWidth.U(value)
686
+ )
687
+ else:
688
+ return False
689
+
690
+ @property
691
+ def IsFloat(self):
692
+ return self._type in (Type.FLOAT, Type.INDIRECT_FLOAT)
693
+
694
+ @property
695
+ def AsFloat(self):
696
+ """Returns current reference as floating point value."""
697
+ if self.IsNull:
698
+ return 0.0
699
+ elif self.IsBool:
700
+ return float(self.AsBool)
701
+ elif self.IsInt:
702
+ return float(self.AsInt)
703
+ elif self._type is Type.FLOAT:
704
+ return _Unpack(F, self._Bytes)
705
+ elif self._type is Type.INDIRECT_FLOAT:
706
+ return _Unpack(F, self._Indirect()[: self._byte_width])
707
+ elif self.IsString:
708
+ return float(self.AsString)
709
+ elif self.IsVector:
710
+ return float(len(self.AsVector))
711
+ elif self.IsTypedVector():
712
+ return float(len(self.AsTypedVector))
713
+ elif self.IsFixedTypedVector():
714
+ return float(len(self.FixedTypedVector))
715
+ else:
716
+ raise self._ConvertError(Type.FLOAT)
717
+
718
+ def MutateFloat(self, value):
719
+ """Mutates underlying floating point value bytes in place.
720
+
721
+ Args:
722
+ value: New float value. It must fit to the byte size of the existing
723
+ encoded value.
724
+
725
+ Returns:
726
+ Whether the value was mutated or not.
727
+ """
728
+ if self._type is Type.FLOAT:
729
+ return _Mutate(
730
+ F, self._buf, value, self._parent_width, BitWidth.B(self._parent_width)
731
+ )
732
+ elif self._type is Type.INDIRECT_FLOAT:
733
+ return _Mutate(
734
+ F,
735
+ self._Indirect(),
736
+ value,
737
+ self._byte_width,
738
+ BitWidth.B(self._byte_width),
739
+ )
740
+ else:
741
+ return False
742
+
743
+ @property
744
+ def IsKey(self):
745
+ return self._type is Type.KEY
746
+
747
+ @property
748
+ def AsKeyBytes(self):
749
+ if self.IsKey:
750
+ return Key(self._Indirect(), self._byte_width).Bytes
751
+ else:
752
+ raise self._ConvertError(Type.KEY)
753
+
754
+ @property
755
+ def AsKey(self):
756
+ if self.IsKey:
757
+ return str(Key(self._Indirect(), self._byte_width))
758
+ else:
759
+ raise self._ConvertError(Type.KEY)
760
+
761
+ @property
762
+ def IsString(self):
763
+ return self._type is Type.STRING
764
+
765
+ @property
766
+ def AsString(self):
767
+ if self.IsString:
768
+ return str(String(self._Indirect(), self._byte_width))
769
+ elif self.IsKey:
770
+ return self.AsKey
771
+ else:
772
+ raise self._ConvertError(Type.STRING)
773
+
774
+ def MutateString(self, value):
775
+ return String(self._Indirect(), self._byte_width).Mutate(value)
776
+
777
+ @property
778
+ def IsBlob(self):
779
+ return self._type is Type.BLOB
780
+
781
+ @property
782
+ def AsBlob(self):
783
+ if self.IsBlob:
784
+ return Blob(self._Indirect(), self._byte_width).Bytes
785
+ else:
786
+ raise self._ConvertError(Type.BLOB)
787
+
788
+ @property
789
+ def IsAnyVector(self):
790
+ return self.IsVector or self.IsTypedVector or self.IsFixedTypedVector()
791
+
792
+ @property
793
+ def IsVector(self):
794
+ return self._type in (Type.VECTOR, Type.MAP)
795
+
796
+ @property
797
+ def AsVector(self):
798
+ if self.IsVector:
799
+ return Vector(self._Indirect(), self._byte_width)
800
+ else:
801
+ raise self._ConvertError(Type.VECTOR)
802
+
803
+ @property
804
+ def IsTypedVector(self):
805
+ return Type.IsTypedVector(self._type)
806
+
807
+ @property
808
+ def AsTypedVector(self):
809
+ if self.IsTypedVector:
810
+ return TypedVector(
811
+ self._Indirect(),
812
+ self._byte_width,
813
+ Type.ToTypedVectorElementType(self._type),
814
+ )
815
+ else:
816
+ raise self._ConvertError("TYPED_VECTOR")
817
+
818
+ @property
819
+ def IsFixedTypedVector(self):
820
+ return Type.IsFixedTypedVector(self._type)
821
+
822
+ @property
823
+ def AsFixedTypedVector(self):
824
+ if self.IsFixedTypedVector:
825
+ element_type, size = Type.ToFixedTypedVectorElementType(self._type)
826
+ return TypedVector(self._Indirect(), self._byte_width, element_type, size)
827
+ else:
828
+ raise self._ConvertError("FIXED_TYPED_VECTOR")
829
+
830
+ @property
831
+ def IsMap(self):
832
+ return self._type is Type.MAP
833
+
834
+ @property
835
+ def AsMap(self):
836
+ if self.IsMap:
837
+ return Map(self._Indirect(), self._byte_width)
838
+ else:
839
+ raise self._ConvertError(Type.MAP)
840
+
841
+ @property
842
+ def Value(self):
843
+ """Converts current reference to value of corresponding type.
844
+
845
+ This is equivalent to calling `AsInt` for integer values, `AsFloat` for
846
+ floating point values, etc.
847
+
848
+ Returns:
849
+ Value of corresponding type.
850
+ """
851
+ if self.IsNull:
852
+ return None
853
+ elif self.IsBool:
854
+ return self.AsBool
855
+ elif self.IsInt:
856
+ return self.AsInt
857
+ elif self.IsFloat:
858
+ return self.AsFloat
859
+ elif self.IsString:
860
+ return self.AsString
861
+ elif self.IsKey:
862
+ return self.AsKey
863
+ elif self.IsBlob:
864
+ return self.AsBlob
865
+ elif self.IsMap:
866
+ return self.AsMap.Value
867
+ elif self.IsVector:
868
+ return self.AsVector.Value
869
+ elif self.IsTypedVector:
870
+ return self.AsTypedVector.Value
871
+ elif self.IsFixedTypedVector:
872
+ return self.AsFixedTypedVector.Value
873
+ else:
874
+ raise TypeError("cannot convert %r to value" % self)
875
+
876
+
877
+ def _IsIterable(obj):
878
+ try:
879
+ iter(obj)
880
+ return True
881
+ except TypeError:
882
+ return False
883
+
884
+
885
+ class Value:
886
+ """Class to represent given value during the encoding process."""
887
+
888
+ @staticmethod
889
+ def Null():
890
+ return Value(0, Type.NULL, BitWidth.W8)
891
+
892
+ @staticmethod
893
+ def Bool(value):
894
+ return Value(value, Type.BOOL, BitWidth.W8)
895
+
896
+ @staticmethod
897
+ def Int(value, bit_width):
898
+ return Value(value, Type.INT, bit_width)
899
+
900
+ @staticmethod
901
+ def UInt(value, bit_width):
902
+ return Value(value, Type.UINT, bit_width)
903
+
904
+ @staticmethod
905
+ def Float(value, bit_width):
906
+ return Value(value, Type.FLOAT, bit_width)
907
+
908
+ @staticmethod
909
+ def Key(offset):
910
+ return Value(offset, Type.KEY, BitWidth.W8)
911
+
912
+ def __init__(self, value, type_, min_bit_width):
913
+ self._value = value
914
+ self._type = type_
915
+
916
+ # For scalars: of itself, for vector: of its elements, for string: length.
917
+ self._min_bit_width = min_bit_width
918
+
919
+ @property
920
+ def Value(self):
921
+ return self._value
922
+
923
+ @property
924
+ def Type(self):
925
+ return self._type
926
+
927
+ @property
928
+ def MinBitWidth(self):
929
+ return self._min_bit_width
930
+
931
+ def StoredPackedType(self, parent_bit_width=BitWidth.W8):
932
+ return Type.Pack(self._type, self.StoredWidth(parent_bit_width))
933
+
934
+ # We have an absolute offset, but want to store a relative offset
935
+ # elem_index elements beyond the current buffer end. Since whether
936
+ # the relative offset fits in a certain byte_width depends on
937
+ # the size of the elements before it (and their alignment), we have
938
+ # to test for each size in turn.
939
+ def ElemWidth(self, buf_size, elem_index=0):
940
+ if Type.IsInline(self._type):
941
+ return self._min_bit_width
942
+ for byte_width in 1, 2, 4, 8:
943
+ offset_loc = (
944
+ buf_size + _PaddingBytes(buf_size, byte_width) + elem_index * byte_width
945
+ )
946
+ bit_width = BitWidth.U(offset_loc - self._value)
947
+ if byte_width == (1 << bit_width):
948
+ return bit_width
949
+ raise ValueError("relative offset is too big")
950
+
951
+ def StoredWidth(self, parent_bit_width=BitWidth.W8):
952
+ if Type.IsInline(self._type):
953
+ return max(self._min_bit_width, parent_bit_width)
954
+ return self._min_bit_width
955
+
956
+ def __repr__(self):
957
+ return "Value(%s, %s, %s)" % (self._value, self._type, self._min_bit_width)
958
+
959
+ def __str__(self):
960
+ return str(self._value)
961
+
962
+
963
+ def InMap(func):
964
+ def wrapper(self, *args, **kwargs):
965
+ if isinstance(args[0], str):
966
+ self.Key(args[0])
967
+ func(self, *args[1:], **kwargs)
968
+ else:
969
+ func(self, *args, **kwargs)
970
+
971
+ return wrapper
972
+
973
+
974
+ def InMapForString(func):
975
+ def wrapper(self, *args):
976
+ if len(args) == 1:
977
+ func(self, args[0])
978
+ elif len(args) == 2:
979
+ self.Key(args[0])
980
+ func(self, args[1])
981
+ else:
982
+ raise ValueError("invalid number of arguments")
983
+
984
+ return wrapper
985
+
986
+
987
+ class Pool:
988
+ """Collection of (data, offset) pairs sorted by data for quick access."""
989
+
990
+ def __init__(self):
991
+ self._pool = [] # sorted list of (data, offset) tuples
992
+
993
+ def FindOrInsert(self, data, offset):
994
+ do = data, offset
995
+ index = _BinarySearch(self._pool, do, lambda a, b: a[0] < b[0])
996
+ if index != -1:
997
+ _, offset = self._pool[index]
998
+ return offset
999
+ self._pool.insert(index, do)
1000
+ return None
1001
+
1002
+ def Clear(self):
1003
+ self._pool = []
1004
+
1005
+ @property
1006
+ def Elements(self):
1007
+ return [data for data, _ in self._pool]
1008
+
1009
+
1010
+ class Builder:
1011
+ """Helper class to encode structural data into flexbuffers format."""
1012
+
1013
+ def __init__(
1014
+ self, share_strings=False, share_keys=True, force_min_bit_width=BitWidth.W8
1015
+ ):
1016
+ self._share_strings = share_strings
1017
+ self._share_keys = share_keys
1018
+ self._force_min_bit_width = force_min_bit_width
1019
+
1020
+ self._string_pool = Pool()
1021
+ self._key_pool = Pool()
1022
+
1023
+ self._finished = False
1024
+ self._buf = bytearray()
1025
+ self._stack = []
1026
+
1027
+ def __len__(self):
1028
+ return len(self._buf)
1029
+
1030
+ @property
1031
+ def StringPool(self):
1032
+ return self._string_pool
1033
+
1034
+ @property
1035
+ def KeyPool(self):
1036
+ return self._key_pool
1037
+
1038
+ def Clear(self):
1039
+ self._string_pool.Clear()
1040
+ self._key_pool.Clear()
1041
+ self._finished = False
1042
+ self._buf = bytearray()
1043
+ self._stack = []
1044
+
1045
+ def Finish(self):
1046
+ """Finishes encoding process and returns underlying buffer."""
1047
+ if self._finished:
1048
+ raise RuntimeError("builder has been already finished")
1049
+
1050
+ # If you hit this exception, you likely have objects that were never
1051
+ # included in a parent. You need to have exactly one root to finish a
1052
+ # buffer. Check your Start/End calls are matched, and all objects are inside
1053
+ # some other object.
1054
+ if len(self._stack) != 1:
1055
+ raise RuntimeError("internal stack size must be one")
1056
+
1057
+ value = self._stack[0]
1058
+ byte_width = self._Align(value.ElemWidth(len(self._buf)))
1059
+ self._WriteAny(value, byte_width=byte_width) # Root value
1060
+ self._Write(U, value.StoredPackedType(), byte_width=1) # Root type
1061
+ self._Write(U, byte_width, byte_width=1) # Root size
1062
+
1063
+ self.finished = True
1064
+ return self._buf
1065
+
1066
+ def _ReadKey(self, offset):
1067
+ key = self._buf[offset:]
1068
+ return key[: key.find(0)]
1069
+
1070
+ def _Align(self, alignment):
1071
+ byte_width = 1 << alignment
1072
+ self._buf.extend(b"\x00" * _PaddingBytes(len(self._buf), byte_width))
1073
+ return byte_width
1074
+
1075
+ def _Write(self, fmt, value, byte_width):
1076
+ self._buf.extend(_Pack(fmt, value, byte_width))
1077
+
1078
+ def _WriteVector(self, fmt, values, byte_width):
1079
+ self._buf.extend(_PackVector(fmt, values, byte_width))
1080
+
1081
+ def _WriteOffset(self, offset, byte_width):
1082
+ relative_offset = len(self._buf) - offset
1083
+ assert byte_width == 8 or relative_offset < (1 << (8 * byte_width))
1084
+ self._Write(U, relative_offset, byte_width)
1085
+
1086
+ def _WriteAny(self, value, byte_width):
1087
+ fmt = {
1088
+ Type.NULL: U,
1089
+ Type.BOOL: U,
1090
+ Type.INT: I,
1091
+ Type.UINT: U,
1092
+ Type.FLOAT: F,
1093
+ }.get(value.Type)
1094
+ if fmt:
1095
+ self._Write(fmt, value.Value, byte_width)
1096
+ else:
1097
+ self._WriteOffset(value.Value, byte_width)
1098
+
1099
+ def _WriteBlob(self, data, append_zero, type_):
1100
+ bit_width = BitWidth.U(len(data))
1101
+ byte_width = self._Align(bit_width)
1102
+ self._Write(U, len(data), byte_width)
1103
+ loc = len(self._buf)
1104
+ self._buf.extend(data)
1105
+ if append_zero:
1106
+ self._buf.append(0)
1107
+ self._stack.append(Value(loc, type_, bit_width))
1108
+ return loc
1109
+
1110
+ def _WriteScalarVector(self, element_type, byte_width, elements, fixed):
1111
+ """Writes scalar vector elements to the underlying buffer."""
1112
+ bit_width = BitWidth.B(byte_width)
1113
+ # If you get this exception, you're trying to write a vector with a size
1114
+ # field that is bigger than the scalars you're trying to write (e.g. a
1115
+ # byte vector > 255 elements). For such types, write a "blob" instead.
1116
+ if BitWidth.U(len(elements)) > bit_width:
1117
+ raise ValueError("too many elements for the given byte_width")
1118
+
1119
+ self._Align(bit_width)
1120
+ if not fixed:
1121
+ self._Write(U, len(elements), byte_width)
1122
+
1123
+ loc = len(self._buf)
1124
+
1125
+ fmt = {Type.INT: I, Type.UINT: U, Type.FLOAT: F}.get(element_type)
1126
+ if not fmt:
1127
+ raise TypeError("unsupported element_type")
1128
+ self._WriteVector(fmt, elements, byte_width)
1129
+
1130
+ type_ = Type.ToTypedVector(element_type, len(elements) if fixed else 0)
1131
+ self._stack.append(Value(loc, type_, bit_width))
1132
+ return loc
1133
+
1134
+ def _CreateVector(self, elements, typed, fixed, keys=None):
1135
+ """Writes vector elements to the underlying buffer."""
1136
+ length = len(elements)
1137
+
1138
+ if fixed and not typed:
1139
+ raise ValueError("fixed vector must be typed")
1140
+
1141
+ # Figure out smallest bit width we can store this vector with.
1142
+ bit_width = max(self._force_min_bit_width, BitWidth.U(length))
1143
+ prefix_elems = 1 # Vector size
1144
+ if keys:
1145
+ bit_width = max(bit_width, keys.ElemWidth(len(self._buf)))
1146
+ prefix_elems += 2 # Offset to the keys vector and its byte width.
1147
+
1148
+ vector_type = Type.KEY
1149
+ # Check bit widths and types for all elements.
1150
+ for i, e in enumerate(elements):
1151
+ bit_width = max(bit_width, e.ElemWidth(len(self._buf), prefix_elems + i))
1152
+
1153
+ if typed:
1154
+ if i == 0:
1155
+ vector_type = e.Type
1156
+ else:
1157
+ if vector_type != e.Type:
1158
+ raise RuntimeError(
1159
+ "typed vector elements must be of the same type"
1160
+ )
1161
+
1162
+ if fixed and not Type.IsFixedTypedVectorElementType(vector_type):
1163
+ raise RuntimeError("must be fixed typed vector element type")
1164
+
1165
+ byte_width = self._Align(bit_width)
1166
+ # Write vector. First the keys width/offset if available, and size.
1167
+ if keys:
1168
+ self._WriteOffset(keys.Value, byte_width)
1169
+ self._Write(U, 1 << keys.MinBitWidth, byte_width)
1170
+
1171
+ if not fixed:
1172
+ self._Write(U, length, byte_width)
1173
+
1174
+ # Then the actual data.
1175
+ loc = len(self._buf)
1176
+ for e in elements:
1177
+ self._WriteAny(e, byte_width)
1178
+
1179
+ # Then the types.
1180
+ if not typed:
1181
+ for e in elements:
1182
+ self._buf.append(e.StoredPackedType(bit_width))
1183
+
1184
+ if keys:
1185
+ type_ = Type.MAP
1186
+ else:
1187
+ if typed:
1188
+ type_ = Type.ToTypedVector(vector_type, length if fixed else 0)
1189
+ else:
1190
+ type_ = Type.VECTOR
1191
+
1192
+ return Value(loc, type_, bit_width)
1193
+
1194
+ def _PushIndirect(self, value, type_, bit_width):
1195
+ byte_width = self._Align(bit_width)
1196
+ loc = len(self._buf)
1197
+ fmt = {Type.INDIRECT_INT: I, Type.INDIRECT_UINT: U, Type.INDIRECT_FLOAT: F}[
1198
+ type_
1199
+ ]
1200
+ self._Write(fmt, value, byte_width)
1201
+ self._stack.append(Value(loc, type_, bit_width))
1202
+
1203
+ @InMapForString
1204
+ def String(self, value):
1205
+ """Encodes string value."""
1206
+ reset_to = len(self._buf)
1207
+ encoded = value.encode("utf-8")
1208
+ loc = self._WriteBlob(encoded, append_zero=True, type_=Type.STRING)
1209
+ if self._share_strings:
1210
+ prev_loc = self._string_pool.FindOrInsert(encoded, loc)
1211
+ if prev_loc is not None:
1212
+ del self._buf[reset_to:]
1213
+ self._stack[-1]._value = loc = prev_loc # pylint: disable=protected-access
1214
+
1215
+ return loc
1216
+
1217
+ @InMap
1218
+ def Blob(self, value):
1219
+ """Encodes binary blob value.
1220
+
1221
+ Args:
1222
+ value: A byte/bytearray value to encode
1223
+
1224
+ Returns:
1225
+ Offset of the encoded value in underlying the byte buffer.
1226
+ """
1227
+ return self._WriteBlob(value, append_zero=False, type_=Type.BLOB)
1228
+
1229
+ def Key(self, value):
1230
+ """Encodes key value.
1231
+
1232
+ Args:
1233
+ value: A byte/bytearray/str value to encode. Byte object must not contain
1234
+ zero bytes. String object must be convertible to ASCII.
1235
+
1236
+ Returns:
1237
+ Offset of the encoded value in the underlying byte buffer.
1238
+ """
1239
+ if isinstance(value, (bytes, bytearray)):
1240
+ encoded = value
1241
+ else:
1242
+ encoded = value.encode("ascii")
1243
+
1244
+ if 0 in encoded:
1245
+ raise ValueError("key contains zero byte")
1246
+
1247
+ loc = len(self._buf)
1248
+ self._buf.extend(encoded)
1249
+ self._buf.append(0)
1250
+ if self._share_keys:
1251
+ prev_loc = self._key_pool.FindOrInsert(encoded, loc)
1252
+ if prev_loc is not None:
1253
+ del self._buf[loc:]
1254
+ loc = prev_loc
1255
+
1256
+ self._stack.append(Value.Key(loc))
1257
+ return loc
1258
+
1259
+ def Null(self, key=None):
1260
+ """Encodes None value."""
1261
+ if key:
1262
+ self.Key(key)
1263
+ self._stack.append(Value.Null())
1264
+
1265
+ @InMap
1266
+ def Bool(self, value):
1267
+ """Encodes boolean value.
1268
+
1269
+ Args:
1270
+ value: A boolean value.
1271
+ """
1272
+ self._stack.append(Value.Bool(value))
1273
+
1274
+ @InMap
1275
+ def Int(self, value, byte_width=0):
1276
+ """Encodes signed integer value.
1277
+
1278
+ Args:
1279
+ value: A signed integer value.
1280
+ byte_width: Number of bytes to use: 1, 2, 4, or 8.
1281
+ """
1282
+ bit_width = BitWidth.I(value) if byte_width == 0 else BitWidth.B(byte_width)
1283
+ self._stack.append(Value.Int(value, bit_width))
1284
+
1285
+ @InMap
1286
+ def IndirectInt(self, value, byte_width=0):
1287
+ """Encodes signed integer value indirectly.
1288
+
1289
+ Args:
1290
+ value: A signed integer value.
1291
+ byte_width: Number of bytes to use: 1, 2, 4, or 8.
1292
+ """
1293
+ bit_width = BitWidth.I(value) if byte_width == 0 else BitWidth.B(byte_width)
1294
+ self._PushIndirect(value, Type.INDIRECT_INT, bit_width)
1295
+
1296
+ @InMap
1297
+ def UInt(self, value, byte_width=0):
1298
+ """Encodes unsigned integer value.
1299
+
1300
+ Args:
1301
+ value: An unsigned integer value.
1302
+ byte_width: Number of bytes to use: 1, 2, 4, or 8.
1303
+ """
1304
+ bit_width = BitWidth.U(value) if byte_width == 0 else BitWidth.B(byte_width)
1305
+ self._stack.append(Value.UInt(value, bit_width))
1306
+
1307
+ @InMap
1308
+ def IndirectUInt(self, value, byte_width=0):
1309
+ """Encodes unsigned integer value indirectly.
1310
+
1311
+ Args:
1312
+ value: An unsigned integer value.
1313
+ byte_width: Number of bytes to use: 1, 2, 4, or 8.
1314
+ """
1315
+ bit_width = BitWidth.U(value) if byte_width == 0 else BitWidth.B(byte_width)
1316
+ self._PushIndirect(value, Type.INDIRECT_UINT, bit_width)
1317
+
1318
+ @InMap
1319
+ def Float(self, value, byte_width=0):
1320
+ """Encodes floating point value.
1321
+
1322
+ Args:
1323
+ value: A floating point value.
1324
+ byte_width: Number of bytes to use: 4 or 8.
1325
+ """
1326
+ bit_width = BitWidth.F(value) if byte_width == 0 else BitWidth.B(byte_width)
1327
+ self._stack.append(Value.Float(value, bit_width))
1328
+
1329
+ @InMap
1330
+ def IndirectFloat(self, value, byte_width=0):
1331
+ """Encodes floating point value indirectly.
1332
+
1333
+ Args:
1334
+ value: A floating point value.
1335
+ byte_width: Number of bytes to use: 4 or 8.
1336
+ """
1337
+ bit_width = BitWidth.F(value) if byte_width == 0 else BitWidth.B(byte_width)
1338
+ self._PushIndirect(value, Type.INDIRECT_FLOAT, bit_width)
1339
+
1340
+ def _StartVector(self):
1341
+ """Starts vector construction."""
1342
+ return len(self._stack)
1343
+
1344
+ def _EndVector(self, start, typed, fixed):
1345
+ """Finishes vector construction by encodung its elements."""
1346
+ vec = self._CreateVector(self._stack[start:], typed, fixed)
1347
+ del self._stack[start:]
1348
+ self._stack.append(vec)
1349
+ return vec.Value
1350
+
1351
+ @contextlib.contextmanager
1352
+ def Vector(self, key=None):
1353
+ if key:
1354
+ self.Key(key)
1355
+
1356
+ try:
1357
+ start = self._StartVector()
1358
+ yield self
1359
+ finally:
1360
+ self._EndVector(start, typed=False, fixed=False)
1361
+
1362
+ @InMap
1363
+ def VectorFromElements(self, elements):
1364
+ """Encodes sequence of any elements as a vector.
1365
+
1366
+ Args:
1367
+ elements: sequence of elements, they may have different types.
1368
+ """
1369
+ with self.Vector():
1370
+ for e in elements:
1371
+ self.Add(e)
1372
+
1373
+ @contextlib.contextmanager
1374
+ def TypedVector(self, key=None):
1375
+ if key:
1376
+ self.Key(key)
1377
+
1378
+ try:
1379
+ start = self._StartVector()
1380
+ yield self
1381
+ finally:
1382
+ self._EndVector(start, typed=True, fixed=False)
1383
+
1384
+ @InMap
1385
+ def TypedVectorFromElements(self, elements, element_type=None):
1386
+ """Encodes sequence of elements of the same type as typed vector.
1387
+
1388
+ Args:
1389
+ elements: Sequence of elements, they must be of the same type.
1390
+ element_type: Suggested element type. Setting it to None means determining
1391
+ correct value automatically based on the given elements.
1392
+ """
1393
+ if isinstance(elements, array.array):
1394
+ if elements.typecode == "f":
1395
+ self._WriteScalarVector(Type.FLOAT, 4, elements, fixed=False)
1396
+ elif elements.typecode == "d":
1397
+ self._WriteScalarVector(Type.FLOAT, 8, elements, fixed=False)
1398
+ elif elements.typecode in ("b", "h", "i", "l", "q"):
1399
+ self._WriteScalarVector(
1400
+ Type.INT, elements.itemsize, elements, fixed=False
1401
+ )
1402
+ elif elements.typecode in ("B", "H", "I", "L", "Q"):
1403
+ self._WriteScalarVector(
1404
+ Type.UINT, elements.itemsize, elements, fixed=False
1405
+ )
1406
+ else:
1407
+ raise ValueError("unsupported array typecode: %s" % elements.typecode)
1408
+ else:
1409
+ add = self.Add if element_type is None else self.Adder(element_type)
1410
+ with self.TypedVector():
1411
+ for e in elements:
1412
+ add(e)
1413
+
1414
+ @InMap
1415
+ def FixedTypedVectorFromElements(self, elements, element_type=None, byte_width=0):
1416
+ """Encodes sequence of elements of the same type as fixed typed vector.
1417
+
1418
+ Args:
1419
+ elements: Sequence of elements, they must be of the same type. Allowed
1420
+ types are `Type.INT`, `Type.UINT`, `Type.FLOAT`. Allowed number of
1421
+ elements are 2, 3, or 4.
1422
+ element_type: Suggested element type. Setting it to None means determining
1423
+ correct value automatically based on the given elements.
1424
+ byte_width: Number of bytes to use per element. For `Type.INT` and
1425
+ `Type.UINT`: 1, 2, 4, or 8. For `Type.FLOAT`: 4 or 8. Setting it to 0
1426
+ means determining correct value automatically based on the given
1427
+ elements.
1428
+ """
1429
+ if not 2 <= len(elements) <= 4:
1430
+ raise ValueError("only 2, 3, or 4 elements are supported")
1431
+
1432
+ types = {type(e) for e in elements}
1433
+ if len(types) != 1:
1434
+ raise TypeError("all elements must be of the same type")
1435
+
1436
+ (type_,) = types
1437
+
1438
+ if element_type is None:
1439
+ element_type = {int: Type.INT, float: Type.FLOAT}.get(type_)
1440
+ if not element_type:
1441
+ raise TypeError("unsupported element_type: %s" % type_)
1442
+
1443
+ if byte_width == 0:
1444
+ width = {
1445
+ Type.UINT: BitWidth.U,
1446
+ Type.INT: BitWidth.I,
1447
+ Type.FLOAT: BitWidth.F,
1448
+ }[element_type]
1449
+ byte_width = 1 << max(width(e) for e in elements)
1450
+
1451
+ self._WriteScalarVector(element_type, byte_width, elements, fixed=True)
1452
+
1453
+ def _StartMap(self):
1454
+ """Starts map construction."""
1455
+ return len(self._stack)
1456
+
1457
+ def _EndMap(self, start):
1458
+ """Finishes map construction by encodung its elements."""
1459
+ # Interleaved keys and values on the stack.
1460
+ stack = self._stack[start:]
1461
+
1462
+ if len(stack) % 2 != 0:
1463
+ raise RuntimeError("must be even number of keys and values")
1464
+
1465
+ for key in stack[::2]:
1466
+ if key.Type is not Type.KEY:
1467
+ raise RuntimeError("all map keys must be of %s type" % Type.KEY)
1468
+
1469
+ pairs = zip(stack[::2], stack[1::2]) # [(key, value), ...]
1470
+ pairs = sorted(pairs, key=lambda pair: self._ReadKey(pair[0].Value))
1471
+
1472
+ del self._stack[start:]
1473
+ for pair in pairs:
1474
+ self._stack.extend(pair)
1475
+
1476
+ keys = self._CreateVector(self._stack[start::2], typed=True, fixed=False)
1477
+ values = self._CreateVector(
1478
+ self._stack[start + 1 :: 2], typed=False, fixed=False, keys=keys
1479
+ )
1480
+
1481
+ del self._stack[start:]
1482
+ self._stack.append(values)
1483
+ return values.Value
1484
+
1485
+ @contextlib.contextmanager
1486
+ def Map(self, key=None):
1487
+ if key:
1488
+ self.Key(key)
1489
+
1490
+ try:
1491
+ start = self._StartMap()
1492
+ yield self
1493
+ finally:
1494
+ self._EndMap(start)
1495
+
1496
+ def MapFromElements(self, elements):
1497
+ start = self._StartMap()
1498
+ for k, v in elements.items():
1499
+ self.Key(k)
1500
+ self.Add(v)
1501
+ self._EndMap(start)
1502
+
1503
+ def Adder(self, type_):
1504
+ return {
1505
+ Type.BOOL: self.Bool,
1506
+ Type.INT: self.Int,
1507
+ Type.INDIRECT_INT: self.IndirectInt,
1508
+ Type.UINT: self.UInt,
1509
+ Type.INDIRECT_UINT: self.IndirectUInt,
1510
+ Type.FLOAT: self.Float,
1511
+ Type.INDIRECT_FLOAT: self.IndirectFloat,
1512
+ Type.KEY: self.Key,
1513
+ Type.BLOB: self.Blob,
1514
+ Type.STRING: self.String,
1515
+ }[type_]
1516
+
1517
+ @InMapForString
1518
+ def Add(self, value):
1519
+ """Encodes value of any supported type."""
1520
+ if value is None:
1521
+ self.Null()
1522
+ elif isinstance(value, bool):
1523
+ self.Bool(value)
1524
+ elif isinstance(value, int):
1525
+ self.Int(value)
1526
+ elif isinstance(value, float):
1527
+ self.Float(value)
1528
+ elif isinstance(value, str):
1529
+ self.String(value)
1530
+ elif isinstance(value, (bytes, bytearray)):
1531
+ self.Blob(value)
1532
+ elif isinstance(value, dict):
1533
+ with self.Map():
1534
+ for k, v in value.items():
1535
+ self.Key(k)
1536
+ self.Add(v)
1537
+ elif isinstance(value, array.array):
1538
+ self.TypedVectorFromElements(value)
1539
+ elif _IsIterable(value):
1540
+ self.VectorFromElements(value)
1541
+ else:
1542
+ raise TypeError("unsupported python type: %s" % type(value))
1543
+
1544
+ @property
1545
+ def LastValue(self):
1546
+ return self._stack[-1]
1547
+
1548
+ @InMap
1549
+ def ReuseValue(self, value):
1550
+ self._stack.append(value)
1551
+
1552
+
1553
+ def GetRoot(buf):
1554
+ """Returns root `Ref` object for the given buffer."""
1555
+ if len(buf) < 3:
1556
+ raise ValueError("buffer is too small")
1557
+ byte_width = buf[-1]
1558
+ return Ref.PackedType(Buf(buf, -(2 + byte_width)), byte_width, packed_type=buf[-2])
1559
+
1560
+
1561
+ def Dumps(obj):
1562
+ """Returns bytearray with the encoded python object."""
1563
+ fbb = Builder()
1564
+ fbb.Add(obj)
1565
+ return fbb.Finish()
1566
+
1567
+
1568
+ def Loads(buf):
1569
+ """Returns python object decoded from the buffer."""
1570
+ return GetRoot(buf).Value