codekingpro/portable-devtools
115k
1# Lint as: python32# 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 at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# 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 and14# limitations under the License.15"""Implementation of FlexBuffers binary format.16 17For more info check https://google.github.io/flatbuffers/flexbuffers.html and18corresponding C++ implementation at19https://github.com/google/flatbuffers/blob/master/include/flatbuffers/flexbuffers.h20"""21 22# pylint: disable=invalid-name23# TODO(dkovalev): Add type hints everywhere, so tools like pytypes could work.24 25import array26import contextlib27import enum28import struct29 30__all__ = ('Type', 'Builder', 'GetRoot', 'Dumps', 'Loads')31 32 33class 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 of37 the elements (and or size field) of the item pointed to (e.g. vector).38 """39 40 W8 = 0 # 2^0 = 1 byte41 W16 = 1 # 2^1 = 2 bytes42 W32 = 2 # 2^2 = 4 bytes43 W64 = 3 # 2^3 = 8 bytes44 45 @staticmethod46 def U(value):47 """Returns the minimum `BitWidth` to encode unsigned integer value."""48 assert value >= 049 50 if value < (1 << 8):51 return BitWidth.W852 elif value < (1 << 16):53 return BitWidth.W1654 elif value < (1 << 32):55 return BitWidth.W3256 elif value < (1 << 64):57 return BitWidth.W6458 else:59 raise ValueError('value is too big to encode: %s' % value)60 61 @staticmethod62 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^n66 # 2 * value < 2^n, when value >= 0 or 2 * (-value) <= 2^n, when value < 067 # 2 * value < 2^n, when value >= 0 or 2 * (-value) - 1 < 2^n, when value < 068 #69 # if value >= 0:70 # return BitWidth.U(2 * value)71 # else:72 # return BitWidth.U(2 * (-value) - 1) # ~x = -x - 173 value *= 274 return BitWidth.U(value if value >= 0 else ~value)75 76 @staticmethod77 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.W3281 return BitWidth.W6482 83 @staticmethod84 def B(byte_width):85 return {1: BitWidth.W8, 2: BitWidth.W16, 4: BitWidth.W32, 8: BitWidth.W64}[86 byte_width87 ]88 89 90I = {1: 'b', 2: 'h', 4: 'i', 8: 'q'} # Integer formats91U = {1: 'B', 2: 'H', 4: 'I', 8: 'Q'} # Unsigned integer formats92F = {4: 'f', 8: 'd'} # Floating point formats93 94 95def _Unpack(fmt, buf):96 return struct.unpack('<%s' % fmt[len(buf)], buf)[0]97 98 99def _UnpackVector(fmt, buf, length):100 byte_width = len(buf) // length101 return struct.unpack('<%d%s' % (length, fmt[byte_width]), buf)102 103 104def _Pack(fmt, value, byte_width):105 return struct.pack('<%s' % fmt[byte_width], value)106 107 108def _PackVector(fmt, values, byte_width):109 return struct.pack('<%d%s' % (len(values), fmt[byte_width]), *values)110 111 112def _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 True116 return False117 118 119# Computes how many bytes you'd have to pad to be able to write an120# "scalar_size" scalar if the buffer had grown to "buf_size",121# "scalar_size" is a power of two.122def _PaddingBytes(buf_size, scalar_size):123 # ((buf_size + (scalar_size - 1)) // scalar_size) * scalar_size - buf_size124 return -buf_size & (scalar_size - 1)125 126 127def _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_bound134def _LowerBound(values, value, pred):135 """Implementation of C++ std::lower_bound() algorithm."""136 first, last = 0, len(values)137 count = last - first138 while count > 0:139 i = first140 step = count // 2141 i += step142 if pred(values[i], value):143 i += 1144 first = i145 count -= step + 1146 else:147 count = step148 return first149 150 151# https://en.cppreference.com/w/cpp/algorithm/binary_search152def _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 index157 return -1158 159 160class 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 actual164 type.165 """166 167 NULL = 0168 INT = 1169 UINT = 2170 FLOAT = 3171 # Types above stored inline, types below store an offset.172 KEY = 4173 STRING = 5174 INDIRECT_INT = 6175 INDIRECT_UINT = 7176 INDIRECT_FLOAT = 8177 MAP = 9178 VECTOR = 10 # Untyped.179 180 VECTOR_INT = 11 # Typed any size (stores no type table).181 VECTOR_UINT = 12182 VECTOR_FLOAT = 13183 VECTOR_KEY = 14184 # DEPRECATED, use VECTOR or VECTOR_KEY instead.185 # Read test.cpp/FlexBuffersDeprecatedTest() for details on why.186 VECTOR_STRING_DEPRECATED = 15187 188 VECTOR_INT2 = 16 # Typed tuple (no type table, no size field).189 VECTOR_UINT2 = 17190 VECTOR_FLOAT2 = 18191 VECTOR_INT3 = 19 # Typed triple (no type table, no size field).192 VECTOR_UINT3 = 20193 VECTOR_FLOAT3 = 21194 VECTOR_INT4 = 22 # Typed quad (no type table, no size field).195 VECTOR_UINT4 = 23196 VECTOR_FLOAT4 = 24197 198 BLOB = 25199 BOOL = 26200 VECTOR_BOOL = 36 # To do the same type of conversion of type to vector type201 202 @staticmethod203 def Pack(type_, bit_width):204 return (int(type_) << 2) | bit_width205 206 @staticmethod207 def Unpack(packed_type):208 return 1 << (packed_type & 0b11), Type(packed_type >> 2)209 210 @staticmethod211 def IsInline(type_):212 return type_ <= Type.FLOAT or type_ == Type.BOOL213 214 @staticmethod215 def IsTypedVector(type_):216 return (217 Type.VECTOR_INT <= type_ <= Type.VECTOR_STRING_DEPRECATED218 or type_ == Type.VECTOR_BOOL219 )220 221 @staticmethod222 def IsTypedVectorElementType(type_):223 return Type.INT <= type_ <= Type.STRING or type_ == Type.BOOL224 225 @staticmethod226 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 @staticmethod233 def IsFixedTypedVector(type_):234 return Type.VECTOR_INT2 <= type_ <= Type.VECTOR_FLOAT4235 236 @staticmethod237 def IsFixedTypedVectorElementType(type_):238 return Type.INT <= type_ <= Type.FLOAT239 240 @staticmethod241 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_INT2247 return Type(fixed_type % 3 + Type.INT), fixed_type // 3 + 2248 249 @staticmethod250 def ToTypedVector(element_type, fixed_len=0):251 """Converts element type to corresponding vector type.252 253 Args:254 element_type: vector element type255 fixed_len: number of elements: 0 for typed vector; 2, 3, or 4 for fixed256 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.INT269 if fixed_len == 0:270 return Type(offset + Type.VECTOR_INT) # TypedVector271 elif fixed_len == 2:272 return Type(offset + Type.VECTOR_INT2) # FixedTypedVector273 elif fixed_len == 3:274 return Type(offset + Type.VECTOR_INT3) # FixedTypedVector275 elif fixed_len == 4:276 return Type(offset + Type.VECTOR_INT4) # FixedTypedVector277 else:278 raise ValueError('unsupported fixed_len: %s' % fixed_len)279 280 281class Buf:282 """Class to access underlying buffer object starting from the given offset."""283 284 def __init__(self, buf, offset):285 self._buf = buf286 self._offset = offset if offset >= 0 else len(buf) + offset287 self._length = len(buf) - self._offset288 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)] = value300 elif isinstance(key, int):301 self._buf[self._offset + key] = key302 else:303 raise TypeError('invalid key type')304 305 def __repr__(self):306 return 'buf[%d:]' % self._offset307 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 321class 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 = buf328 self._byte_width = byte_width329 330 @property331 def ByteWidth(self):332 return self._byte_width333 334 335class 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 = size346 347 @property348 def SizeBytes(self):349 return self._buf[-self._byte_width : 0]350 351 def __len__(self):352 return self._size353 354 355class Blob(Sized):356 """Data accessor for the encoded blob bytes."""357 358 __slots__ = ()359 360 @property361 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 368class String(Sized):369 """Data accessor for the encoded string bytes."""370 371 __slots__ = ()372 373 @property374 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 less382 or equal UTF-8-encoded bytes than the existing one to successfully383 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] = encoded393 self._buf[n : len(self)] = bytearray(len(self) - n)394 return True395 return False396 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 404class 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 == 1411 super().__init__(buf, byte_width)412 413 @property414 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 427class 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 @property443 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 455class 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-width465 # of the size field, see the declaration of466 # 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 with469 # embedded nulls.470 element_type = Type.KEY471 472 self._element_type = element_type473 474 @property475 def Bytes(self):476 return self._buf[: self._byte_width * len(self)]477 478 @property479 def ElementType(self):480 return self._element_type481 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 @property492 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 521class Map(Vector):522 """Data accessor for the encoded map bytes."""523 524 @staticmethod525 def CompareKeys(a, b):526 if isinstance(a, Ref):527 a = a.AsKeyBytes528 if isinstance(b, Ref):529 b = b.AsKeyBytes530 return a < b531 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 @property543 def Keys(self):544 byte_width = _Unpack(545 U, self._buf[-2 * self._byte_width : -self._byte_width]546 )547 buf = self._buf.Indirect(-3 * self._byte_width, self._byte_width)548 return TypedVector(buf, byte_width, Type.KEY)549 550 @property551 def Values(self):552 return Vector(self._buf, self._byte_width)553 554 @property555 def Value(self):556 return {k.Value: v.Value for k, v in zip(self.Keys, self.Values)}557 558 def __repr__(self):559 return 'Map(%s, size=%d)' % (self._buf, len(self))560 561 562class Ref:563 """Data accessor for the encoded data bytes."""564 565 __slots__ = '_buf', '_parent_width', '_byte_width', '_type'566 567 @staticmethod568 def PackedType(buf, parent_width, packed_type):569 byte_width, type_ = Type.Unpack(packed_type)570 return Ref(buf, parent_width, byte_width, type_)571 572 def __init__(self, buf, parent_width, byte_width, type_):573 self._buf = buf574 self._parent_width = parent_width575 self._byte_width = byte_width576 self._type = type_577 578 def __repr__(self):579 return 'Ref(%s, parent_width=%d, byte_width=%d, type_=%s)' % (580 self._buf,581 self._parent_width,582 self._byte_width,583 self._type,584 )585 586 @property587 def _Bytes(self):588 return self._buf[: self._parent_width]589 590 def _ConvertError(self, target_type):591 raise TypeError('cannot convert %s to %s' % (self._type, target_type))592 593 def _Indirect(self):594 return self._buf.Indirect(0, self._parent_width)595 596 @property597 def IsNull(self):598 return self._type is Type.NULL599 600 @property601 def IsBool(self):602 return self._type is Type.BOOL603 604 @property605 def AsBool(self):606 if self._type is Type.BOOL:607 return bool(_Unpack(U, self._Bytes))608 else:609 return self.AsInt != 0610 611 def MutateBool(self, value):612 """Mutates underlying boolean value bytes in place.613 614 Args:615 value: New boolean value.616 617 Returns:618 Whether the value was mutated or not.619 """620 return self.IsBool and _Mutate(621 U, self._buf, value, self._parent_width, BitWidth.W8622 )623 624 @property625 def IsNumeric(self):626 return self.IsInt or self.IsFloat627 628 @property629 def IsInt(self):630 return self._type in (631 Type.INT,632 Type.INDIRECT_INT,633 Type.UINT,634 Type.INDIRECT_UINT,635 )636 637 @property638 def AsInt(self):639 """Returns current reference as integer value."""640 if self.IsNull:641 return 0642 elif self.IsBool:643 return int(self.AsBool)644 elif self._type is Type.INT:645 return _Unpack(I, self._Bytes)646 elif self._type is Type.INDIRECT_INT:647 return _Unpack(I, self._Indirect()[: self._byte_width])648 if self._type is Type.UINT:649 return _Unpack(U, self._Bytes)650 elif self._type is Type.INDIRECT_UINT:651 return _Unpack(U, self._Indirect()[: self._byte_width])652 elif self.IsString:653 return len(self.AsString)654 elif self.IsKey:655 return len(self.AsKey)656 elif self.IsBlob:657 return len(self.AsBlob)658 elif self.IsVector:659 return len(self.AsVector)660 elif self.IsTypedVector:661 return len(self.AsTypedVector)662 elif self.IsFixedTypedVector:663 return len(self.AsFixedTypedVector)664 else:665 raise self._ConvertError(Type.INT)666 667 def MutateInt(self, value):668 """Mutates underlying integer value bytes in place.669 670 Args:671 value: New integer value. It must fit to the byte size of the existing672 encoded value.673 674 Returns:675 Whether the value was mutated or not.676 """677 if self._type is Type.INT:678 return _Mutate(I, self._buf, value, self._parent_width, BitWidth.I(value))679 elif self._type is Type.INDIRECT_INT:680 return _Mutate(681 I, self._Indirect(), value, self._byte_width, BitWidth.I(value)682 )683 elif self._type is Type.UINT:684 return _Mutate(U, self._buf, value, self._parent_width, BitWidth.U(value))685 elif self._type is Type.INDIRECT_UINT:686 return _Mutate(687 U, self._Indirect(), value, self._byte_width, BitWidth.U(value)688 )689 else:690 return False691 692 @property693 def IsFloat(self):694 return self._type in (Type.FLOAT, Type.INDIRECT_FLOAT)695 696 @property697 def AsFloat(self):698 """Returns current reference as floating point value."""699 if self.IsNull:700 return 0.0701 elif self.IsBool:702 return float(self.AsBool)703 elif self.IsInt:704 return float(self.AsInt)705 elif self._type is Type.FLOAT:706 return _Unpack(F, self._Bytes)707 elif self._type is Type.INDIRECT_FLOAT:708 return _Unpack(F, self._Indirect()[: self._byte_width])709 elif self.IsString:710 return float(self.AsString)711 elif self.IsVector:712 return float(len(self.AsVector))713 elif self.IsTypedVector():714 return float(len(self.AsTypedVector))715 elif self.IsFixedTypedVector():716 return float(len(self.FixedTypedVector))717 else:718 raise self._ConvertError(Type.FLOAT)719 720 def MutateFloat(self, value):721 """Mutates underlying floating point value bytes in place.722 723 Args:724 value: New float value. It must fit to the byte size of the existing725 encoded value.726 727 Returns:728 Whether the value was mutated or not.729 """730 if self._type is Type.FLOAT:731 return _Mutate(732 F,733 self._buf,734 value,735 self._parent_width,736 BitWidth.B(self._parent_width),737 )738 elif self._type is Type.INDIRECT_FLOAT:739 return _Mutate(740 F,741 self._Indirect(),742 value,743 self._byte_width,744 BitWidth.B(self._byte_width),745 )746 else:747 return False748 749 @property750 def IsKey(self):751 return self._type is Type.KEY752 753 @property754 def AsKeyBytes(self):755 if self.IsKey:756 return Key(self._Indirect(), self._byte_width).Bytes757 else:758 raise self._ConvertError(Type.KEY)759 760 @property761 def AsKey(self):762 if self.IsKey:763 return str(Key(self._Indirect(), self._byte_width))764 else:765 raise self._ConvertError(Type.KEY)766 767 @property768 def IsString(self):769 return self._type is Type.STRING770 771 @property772 def AsStringBytes(self):773 if self.IsString:774 return String(self._Indirect(), self._byte_width).Bytes775 elif self.IsKey:776 return self.AsKeyBytes777 else:778 raise self._ConvertError(Type.STRING)779 780 @property781 def AsString(self):782 if self.IsString:783 return str(String(self._Indirect(), self._byte_width))784 elif self.IsKey:785 return self.AsKey786 else:787 raise self._ConvertError(Type.STRING)788 789 def MutateString(self, value):790 return String(self._Indirect(), self._byte_width).Mutate(value)791 792 @property793 def IsBlob(self):794 return self._type is Type.BLOB795 796 @property797 def AsBlob(self):798 if self.IsBlob:799 return Blob(self._Indirect(), self._byte_width).Bytes800 else:801 raise self._ConvertError(Type.BLOB)802 803 @property804 def IsAnyVector(self):805 return self.IsVector or self.IsTypedVector or self.IsFixedTypedVector()806 807 @property808 def IsVector(self):809 return self._type in (Type.VECTOR, Type.MAP)810 811 @property812 def AsVector(self):813 if self.IsVector:814 return Vector(self._Indirect(), self._byte_width)815 else:816 raise self._ConvertError(Type.VECTOR)817 818 @property819 def IsTypedVector(self):820 return Type.IsTypedVector(self._type)821 822 @property823 def AsTypedVector(self):824 if self.IsTypedVector:825 return TypedVector(826 self._Indirect(),827 self._byte_width,828 Type.ToTypedVectorElementType(self._type),829 )830 else:831 raise self._ConvertError('TYPED_VECTOR')832 833 @property834 def IsFixedTypedVector(self):835 return Type.IsFixedTypedVector(self._type)836 837 @property838 def AsFixedTypedVector(self):839 if self.IsFixedTypedVector:840 element_type, size = Type.ToFixedTypedVectorElementType(self._type)841 return TypedVector(self._Indirect(), self._byte_width, element_type, size)842 else:843 raise self._ConvertError('FIXED_TYPED_VECTOR')844 845 @property846 def IsMap(self):847 return self._type is Type.MAP848 849 @property850 def AsMap(self):851 if self.IsMap:852 return Map(self._Indirect(), self._byte_width)853 else:854 raise self._ConvertError(Type.MAP)855 856 @property857 def Value(self):858 """Converts current reference to value of corresponding type.859 860 This is equivalent to calling `AsInt` for integer values, `AsFloat` for861 floating point values, etc.862 863 Returns:864 Value of corresponding type.865 """866 if self.IsNull:867 return None868 elif self.IsBool:869 return self.AsBool870 elif self.IsInt:871 return self.AsInt872 elif self.IsFloat:873 return self.AsFloat874 elif self.IsString:875 return self.AsString876 elif self.IsKey:877 return self.AsKey878 elif self.IsBlob:879 return self.AsBlob880 elif self.IsMap:881 return self.AsMap.Value882 elif self.IsVector:883 return self.AsVector.Value884 elif self.IsTypedVector:885 return self.AsTypedVector.Value886 elif self.IsFixedTypedVector:887 return self.AsFixedTypedVector.Value888 else:889 raise TypeError('cannot convert %r to value' % self)890 891 892def _IsIterable(obj):893 try:894 iter(obj)895 return True896 except TypeError:897 return False898 899 900class Value:901 """Class to represent given value during the encoding process."""902 903 @staticmethod904 def Null():905 return Value(0, Type.NULL, BitWidth.W8)906 907 @staticmethod908 def Bool(value):909 return Value(value, Type.BOOL, BitWidth.W8)910 911 @staticmethod912 def Int(value, bit_width):913 return Value(value, Type.INT, bit_width)914 915 @staticmethod916 def UInt(value, bit_width):917 return Value(value, Type.UINT, bit_width)918 919 @staticmethod920 def Float(value, bit_width):921 return Value(value, Type.FLOAT, bit_width)922 923 @staticmethod924 def Key(offset):925 return Value(offset, Type.KEY, BitWidth.W8)926 927 def __init__(self, value, type_, min_bit_width):928 self._value = value929 self._type = type_930 931 # For scalars: of itself, for vector: of its elements, for string: length.932 self._min_bit_width = min_bit_width933 934 @property935 def Value(self):936 return self._value937 938 @property939 def Type(self):940 return self._type941 942 @property943 def MinBitWidth(self):944 return self._min_bit_width945 946 def StoredPackedType(self, parent_bit_width=BitWidth.W8):947 return Type.Pack(self._type, self.StoredWidth(parent_bit_width))948 949 # We have an absolute offset, but want to store a relative offset950 # elem_index elements beyond the current buffer end. Since whether951 # the relative offset fits in a certain byte_width depends on952 # the size of the elements before it (and their alignment), we have953 # to test for each size in turn.954 def ElemWidth(self, buf_size, elem_index=0):955 if Type.IsInline(self._type):956 return self._min_bit_width957 for byte_width in 1, 2, 4, 8:958 offset_loc = (959 buf_size960 + _PaddingBytes(buf_size, byte_width)961 + elem_index * byte_width962 )963 bit_width = BitWidth.U(offset_loc - self._value)964 if byte_width == (1 << bit_width):965 return bit_width966 raise ValueError('relative offset is too big')967 968 def StoredWidth(self, parent_bit_width=BitWidth.W8):969 if Type.IsInline(self._type):970 return max(self._min_bit_width, parent_bit_width)971 return self._min_bit_width972 973 def __repr__(self):974 return 'Value(%s, %s, %s)' % (self._value, self._type, self._min_bit_width)975 976 def __str__(self):977 return str(self._value)978 979 980def InMap(func):981 def wrapper(self, *args, **kwargs):982 if isinstance(args[0], str):983 self.Key(args[0])984 func(self, *args[1:], **kwargs)985 else:986 func(self, *args, **kwargs)987 988 return wrapper989 990 991def InMapForString(func):992 def wrapper(self, *args):993 if len(args) == 1:994 func(self, args[0])995 elif len(args) == 2:996 self.Key(args[0])997 func(self, args[1])998 else:999 raise ValueError('invalid number of arguments')1000 1001 return wrapper1002 1003 1004class Pool:1005 """Collection of (data, offset) pairs sorted by data for quick access."""1006 1007 def __init__(self):1008 self._pool = [] # sorted list of (data, offset) tuples1009 1010 def FindOrInsert(self, data, offset):1011 do = data, offset1012 index = _BinarySearch(self._pool, do, lambda a, b: a[0] < b[0])1013 if index != -1:1014 _, offset = self._pool[index]1015 return offset1016 self._pool.insert(index, do)1017 return None1018 1019 def Clear(self):1020 self._pool = []1021 1022 @property1023 def Elements(self):1024 return [data for data, _ in self._pool]1025 1026 1027class Builder:1028 """Helper class to encode structural data into flexbuffers format."""1029 1030 def __init__(1031 self,1032 share_strings=False,1033 share_keys=True,1034 force_min_bit_width=BitWidth.W8,1035 ):1036 self._share_strings = share_strings1037 self._share_keys = share_keys1038 self._force_min_bit_width = force_min_bit_width1039 1040 self._string_pool = Pool()1041 self._key_pool = Pool()1042 1043 self._finished = False1044 self._buf = bytearray()1045 self._stack = []1046 1047 def __len__(self):1048 return len(self._buf)1049 1050 @property1051 def StringPool(self):1052 return self._string_pool1053 1054 @property1055 def KeyPool(self):1056 return self._key_pool1057 1058 def Clear(self):1059 self._string_pool.Clear()1060 self._key_pool.Clear()1061 self._finished = False1062 self._buf = bytearray()1063 self._stack = []1064 1065 def Finish(self):1066 """Finishes encoding process and returns underlying buffer."""1067 if self._finished:1068 raise RuntimeError('builder has been already finished')1069 1070 # If you hit this exception, you likely have objects that were never1071 # included in a parent. You need to have exactly one root to finish a1072 # buffer. Check your Start/End calls are matched, and all objects are inside1073 # some other object.1074 if len(self._stack) != 1:1075 raise RuntimeError('internal stack size must be one')1076 1077 value = self._stack[0]1078 byte_width = self._Align(value.ElemWidth(len(self._buf)))1079 self._WriteAny(value, byte_width=byte_width) # Root value1080 self._Write(U, value.StoredPackedType(), byte_width=1) # Root type1081 self._Write(U, byte_width, byte_width=1) # Root size1082 1083 self.finished = True1084 return self._buf1085 1086 def _ReadKey(self, offset):1087 key = self._buf[offset:]1088 return key[: key.find(0)]1089 1090 def _Align(self, alignment):1091 byte_width = 1 << alignment1092 self._buf.extend(b'\x00' * _PaddingBytes(len(self._buf), byte_width))1093 return byte_width1094 1095 def _Write(self, fmt, value, byte_width):1096 self._buf.extend(_Pack(fmt, value, byte_width))1097 1098 def _WriteVector(self, fmt, values, byte_width):1099 self._buf.extend(_PackVector(fmt, values, byte_width))1100 1101 def _WriteOffset(self, offset, byte_width):1102 relative_offset = len(self._buf) - offset1103 assert byte_width == 8 or relative_offset < (1 << (8 * byte_width))1104 self._Write(U, relative_offset, byte_width)1105 1106 def _WriteAny(self, value, byte_width):1107 fmt = {1108 Type.NULL: U,1109 Type.BOOL: U,1110 Type.INT: I,1111 Type.UINT: U,1112 Type.FLOAT: F,1113 }.get(value.Type)1114 if fmt:1115 self._Write(fmt, value.Value, byte_width)1116 else:1117 self._WriteOffset(value.Value, byte_width)1118 1119 def _WriteBlob(self, data, append_zero, type_):1120 bit_width = BitWidth.U(len(data))1121 byte_width = self._Align(bit_width)1122 self._Write(U, len(data), byte_width)1123 loc = len(self._buf)1124 self._buf.extend(data)1125 if append_zero:1126 self._buf.append(0)1127 self._stack.append(Value(loc, type_, bit_width))1128 return loc1129 1130 def _WriteScalarVector(self, element_type, byte_width, elements, fixed):1131 """Writes scalar vector elements to the underlying buffer."""1132 bit_width = BitWidth.B(byte_width)1133 # If you get this exception, you're trying to write a vector with a size1134 # field that is bigger than the scalars you're trying to write (e.g. a1135 # byte vector > 255 elements). For such types, write a "blob" instead.1136 if BitWidth.U(len(elements)) > bit_width:1137 raise ValueError('too many elements for the given byte_width')1138 1139 self._Align(bit_width)1140 if not fixed:1141 self._Write(U, len(elements), byte_width)1142 1143 loc = len(self._buf)1144 1145 fmt = {Type.INT: I, Type.UINT: U, Type.FLOAT: F}.get(element_type)1146 if not fmt:1147 raise TypeError('unsupported element_type')1148 self._WriteVector(fmt, elements, byte_width)1149 1150 type_ = Type.ToTypedVector(element_type, len(elements) if fixed else 0)1151 self._stack.append(Value(loc, type_, bit_width))1152 return loc1153 1154 def _CreateVector(self, elements, typed, fixed, keys=None):1155 """Writes vector elements to the underlying buffer."""1156 length = len(elements)1157 1158 if fixed and not typed:1159 raise ValueError('fixed vector must be typed')1160 1161 # Figure out smallest bit width we can store this vector with.1162 bit_width = max(self._force_min_bit_width, BitWidth.U(length))1163 prefix_elems = 1 # Vector size1164 if keys:1165 bit_width = max(bit_width, keys.ElemWidth(len(self._buf)))1166 prefix_elems += 2 # Offset to the keys vector and its byte width.1167 1168 vector_type = Type.KEY1169 # Check bit widths and types for all elements.1170 for i, e in enumerate(elements):1171 bit_width = max(bit_width, e.ElemWidth(len(self._buf), prefix_elems + i))1172 1173 if typed:1174 if i == 0:1175 vector_type = e.Type1176 else:1177 if vector_type != e.Type:1178 raise RuntimeError('typed vector elements must be of the same type')1179 1180 if fixed and not Type.IsFixedTypedVectorElementType(vector_type):1181 raise RuntimeError('must be fixed typed vector element type')1182 1183 byte_width = self._Align(bit_width)1184 # Write vector. First the keys width/offset if available, and size.1185 if keys:1186 self._WriteOffset(keys.Value, byte_width)1187 self._Write(U, 1 << keys.MinBitWidth, byte_width)1188 1189 if not fixed:1190 self._Write(U, length, byte_width)1191 1192 # Then the actual data.1193 loc = len(self._buf)1194 for e in elements:1195 self._WriteAny(e, byte_width)1196 1197 # Then the types.1198 if not typed:1199 for e in elements:1200 self._buf.append(e.StoredPackedType(bit_width))