2020from . import packer
2121from .compat import memoryview_type
2222from .compat import NumpyRequiredForThisFeature , import_numpy
23- from .compat import range_func
2423from .number_types import (SOffsetTFlags , UOffsetTFlags , VOffsetTFlags )
2524
2625np = import_numpy ()
@@ -159,7 +158,7 @@ def __init__(self, initialSize=1024):
159158 self .vtables = {}
160159 self .nested = False
161160 self .forceDefaults = False
162- self .sharedStrings = {}
161+ self .sharedStrings = None
163162 ## @endcond
164163 self .finished = False
165164
@@ -172,7 +171,7 @@ def Clear(self) -> None:
172171 self .vtables = {}
173172 self .nested = False
174173 self .forceDefaults = False
175- self .sharedStrings = {}
174+ self .sharedStrings = None
176175 self .vectorNumElems = None
177176 ## @endcond
178177 self .finished = False
@@ -192,7 +191,7 @@ def Output(self):
192191 if not self .finished :
193192 raise BuilderNotFinishedError ()
194193
195- return self .Bytes [self .Head () :]
194+ return self .Bytes [self .head :]
196195
197196 ## @cond FLATBUFFERS_INTERNAL
198197 def StartObject (self , numfields ):
@@ -201,7 +200,7 @@ def StartObject(self, numfields):
201200 self .assertNotNested ()
202201
203202 # use 32-bit offsets so that arithmetic doesn't overflow.
204- self .current_vtable = [0 for _ in range_func ( numfields )]
203+ self .current_vtable = [0 ] * numfields
205204 self .objectEnd = self .Offset ()
206205 self .nested = True
207206
@@ -255,6 +254,7 @@ def WriteVtable(self):
255254 i = len (self .current_vtable ) - 1
256255 trailing = 0
257256 trim = True
257+ vt_entries = []
258258 while i >= 0 :
259259 off = 0
260260 elem = self .current_vtable [i ]
@@ -270,18 +270,20 @@ def WriteVtable(self):
270270 off = objectOffset - elem
271271 trim = False
272272
273- self . PrependVOffsetT (off )
273+ vt_entries . append (off )
274274
275275 # The two metadata fields are written last.
276276
277277 # First, store the object bytesize:
278278 objectSize = UOffsetTFlags .py_type (objectOffset - self .objectEnd )
279- self . PrependVOffsetT (VOffsetTFlags .py_type (objectSize ))
279+ vt_entries . append (VOffsetTFlags .py_type (objectSize ))
280280
281281 # Second, store the vtable bytesize:
282282 vBytes = len (self .current_vtable ) - trailing + VtableMetadataFields
283283 vBytes *= N .VOffsetTFlags .bytewidth
284- self .PrependVOffsetT (VOffsetTFlags .py_type (vBytes ))
284+ vt_entries .append (VOffsetTFlags .py_type (vBytes ))
285+
286+ self .WriteVtableEntries (vt_entries )
285287
286288 # Next, write the offset to the new vtable in the
287289 # already-allocated SOffsetT at the beginning of this object:
@@ -306,20 +308,31 @@ def WriteVtable(self):
306308 encode .Write (
307309 packer .soffset ,
308310 self .Bytes ,
309- self .Head () ,
311+ self .head ,
310312 SOffsetTFlags .py_type (vt2Offset - objectOffset ),
311313 )
312314
313315 self .current_vtable = None
314316 return objectOffset
315317
318+ def WriteVtableEntries (self , entries ):
319+ """Write a contiguous block of VOffsetT values with a single prep call."""
320+ count = len (entries )
321+ if count == 0 :
322+ return
323+ elem_size = N .VOffsetTFlags .bytewidth
324+ total_bytes = elem_size * count
325+ self .Prep (N .VOffsetTFlags .bytewidth , total_bytes - elem_size )
326+ for value in entries :
327+ self .PlaceVOffsetT (value )
328+
316329 def EndObject (self ):
317330 """EndObject writes data necessary to finish object construction."""
318331 self .assertNested ()
319332 self .nested = False
320333 return self .WriteVtable ()
321334
322- def growByteBuffer (self ):
335+ def GrowByteBuffer (self ):
323336 """Doubles the size of the byteslice, and copies the old data towards
324337
325338 the end of the new buffer (since we build the buffer backwards).
@@ -350,12 +363,15 @@ def Head(self):
350363 ## @cond FLATBUFFERS_INTERNAL
351364 def Offset (self ):
352365 """Offset relative to the end of the buffer."""
353- return UOffsetTFlags .py_type (len (self .Bytes ) - self .Head () )
366+ return UOffsetTFlags .py_type (len (self .Bytes ) - self .head )
354367
355368 def Pad (self , n ):
356369 """Pad places zeros at the current offset."""
357- for i in range_func (n ):
358- self .Place (0 , N .Uint8Flags )
370+ if n <= 0 :
371+ return
372+ new_head = self .head - n
373+ self .Bytes [new_head : self .head ] = b"\x00 " * n
374+ self .head = UOffsetTFlags .py_type (new_head )
359375
360376 def Prep (self , size , additionalBytes ):
361377 """Prep prepares to write an element of `size` after `additional_bytes`
@@ -372,15 +388,19 @@ def Prep(self, size, additionalBytes):
372388
373389 # Find the amount of alignment needed such that `size` is properly
374390 # aligned after `additionalBytes`:
375- alignSize = (~ (len (self .Bytes ) - self .Head () + additionalBytes )) + 1
391+ head = self .head
392+ buf_len = len (self .Bytes )
393+ alignSize = (~ (buf_len - head + additionalBytes )) + 1
376394 alignSize &= size - 1
377395
378396 # Reallocate the buffer if needed:
379- while self .Head () < alignSize + size + additionalBytes :
380- oldBufSize = len (self .Bytes )
381- self .growByteBuffer ()
382- updated_head = self .head + len (self .Bytes ) - oldBufSize
383- self .head = UOffsetTFlags .py_type (updated_head )
397+ needed = alignSize + size + additionalBytes
398+ while head < needed :
399+ oldBufSize = buf_len
400+ self .GrowByteBuffer ()
401+ buf_len = len (self .Bytes )
402+ head += buf_len - oldBufSize
403+ self .head = UOffsetTFlags .py_type (head )
384404 self .Pad (alignSize )
385405
386406 def PrependSOffsetTRelative (self , off ):
@@ -455,7 +475,9 @@ def CreateSharedString(self, s, encoding="utf-8", errors="strict"):
455475 before calling CreateString.
456476 """
457477
458- if s in self .sharedStrings :
478+ if not self .sharedStrings :
479+ self .sharedStrings = {}
480+ elif s in self .sharedStrings :
459481 return self .sharedStrings [s ]
460482
461483 off = self .CreateString (s , encoding , errors )
@@ -478,16 +500,17 @@ def CreateString(self, s, encoding="utf-8", errors="strict"):
478500 else :
479501 raise TypeError ("non-string passed to CreateString" )
480502
481- self .Prep (N .UOffsetTFlags .bytewidth , (len (x ) + 1 ) * N .Uint8Flags .bytewidth )
503+ payload_len = len (x )
504+ self .Prep (
505+ N .UOffsetTFlags .bytewidth , (payload_len + 1 ) * N .Uint8Flags .bytewidth
506+ )
482507 self .Place (0 , N .Uint8Flags )
483508
484- l = UOffsetTFlags .py_type (len (s ))
485- ## @cond FLATBUFFERS_INTERNAL
486- self .head = UOffsetTFlags .py_type (self .Head () - l )
487- ## @endcond
488- self .Bytes [self .Head () : self .Head () + l ] = x
509+ new_head = self .head - payload_len
510+ self .head = UOffsetTFlags .py_type (new_head )
511+ self .Bytes [new_head : new_head + payload_len ] = x
489512
490- self .vectorNumElems = len ( x )
513+ self .vectorNumElems = payload_len
491514 return self .EndVector ()
492515
493516 def CreateByteVector (self , x ):
@@ -501,15 +524,13 @@ def CreateByteVector(self, x):
501524 if not isinstance (x , compat .binary_types ):
502525 raise TypeError ("non-byte vector passed to CreateByteVector" )
503526
504- self .Prep (N .UOffsetTFlags .bytewidth , len (x ) * N .Uint8Flags .bytewidth )
505-
506- l = UOffsetTFlags .py_type (len (x ))
507- ## @cond FLATBUFFERS_INTERNAL
508- self .head = UOffsetTFlags .py_type (self .Head () - l )
509- ## @endcond
510- self .Bytes [self .Head () : self .Head () + l ] = x
527+ data_len = len (x )
528+ self .Prep (N .UOffsetTFlags .bytewidth , data_len * N .Uint8Flags .bytewidth )
529+ new_head = self .head - data_len
530+ self .head = UOffsetTFlags .py_type (new_head )
531+ self .Bytes [new_head : new_head + data_len ] = x
511532
512- self .vectorNumElems = len ( x )
533+ self .vectorNumElems = data_len
513534 return self .EndVector ()
514535
515536 def CreateNumpyVector (self , x ):
@@ -537,13 +558,11 @@ def CreateNumpyVector(self, x):
537558 x_lend = x .byteswap (inplace = False )
538559
539560 # Calculate total length
540- l = UOffsetTFlags .py_type (x_lend .itemsize * x_lend .size )
541- ## @cond FLATBUFFERS_INTERNAL
542- self .head = UOffsetTFlags .py_type (self .Head () - l )
543- ## @endcond
544-
545- # tobytes ensures c_contiguous ordering
546- self .Bytes [self .Head () : self .Head () + l ] = x_lend .tobytes (order = "C" )
561+ payload = x_lend .tobytes (order = "C" )
562+ payload_len = len (payload )
563+ new_head = self .head - payload_len
564+ self .head = UOffsetTFlags .py_type (new_head )
565+ self .Bytes [new_head : new_head + payload_len ] = payload
547566
548567 self .vectorNumElems = x .size
549568 return self .EndVector ()
@@ -613,11 +632,11 @@ def __Finish(self, rootTable, sizePrefix, file_identifier=None):
613632
614633 self .PrependUOffsetTRelative (rootTable )
615634 if sizePrefix :
616- size = len (self .Bytes ) - self .Head ()
635+ size = len (self .Bytes ) - self .head
617636 N .enforce_number (size , N .Int32Flags )
618637 self .PrependInt32 (size )
619638 self .finished = True
620- return self .Head ()
639+ return self .head
621640
622641 def Finish (self , rootTable , file_identifier = None ):
623642 """Finish finalizes a buffer, pointing to the given `rootTable`."""
@@ -632,8 +651,31 @@ def FinishSizePrefixed(self, rootTable, file_identifier=None):
632651
633652 ## @cond FLATBUFFERS_INTERNAL
634653 def Prepend (self , flags , off ):
635- self .Prep (flags .bytewidth , 0 )
636- self .Place (off , flags )
654+ size = flags .bytewidth
655+ if size > self .minalign :
656+ self .minalign = size
657+
658+ head = self .head
659+ buf_len = len (self .Bytes )
660+ alignSize = (~ (buf_len - head )) + 1
661+ alignSize &= size - 1
662+
663+ needed = alignSize + size
664+ while head < needed :
665+ oldBufSize = buf_len
666+ self .GrowByteBuffer ()
667+ buf_len = len (self .Bytes )
668+ head += buf_len - oldBufSize
669+
670+ if alignSize :
671+ new_head = head - alignSize
672+ self .Bytes [new_head :head ] = b"\x00 " * alignSize
673+ head = new_head
674+
675+ N .enforce_number (off , flags )
676+ head -= size
677+ self .head = UOffsetTFlags .py_type (head )
678+ encode .Write (flags .packer_type , self .Bytes , head , off )
637679
638680 def PrependSlot (self , flags , o , x , d ):
639681 if x is not None :
@@ -800,46 +842,46 @@ def ForceDefaults(self, forceDefaults):
800842
801843 ##############################################################
802844
803- ## @cond FLATBUFFERS_INTERNAL
804- def PrependVOffsetT (self , x ):
805- self .Prepend (N .VOffsetTFlags , x )
806-
807845 def Place (self , x , flags ):
808846 """Place prepends a value specified by `flags` to the Builder,
809847
810848 without checking for available space.
811849 """
812850
813851 N .enforce_number (x , flags )
814- self .head = self .head - flags .bytewidth
815- encode .Write (flags .packer_type , self .Bytes , self .Head (), x )
852+ new_head = self .head - flags .bytewidth
853+ self .head = UOffsetTFlags .py_type (new_head )
854+ encode .Write (flags .packer_type , self .Bytes , new_head , x )
816855
817856 def PlaceVOffsetT (self , x ):
818857 """PlaceVOffsetT prepends a VOffsetT to the Builder, without checking
819858
820859 for space.
821860 """
822861 N .enforce_number (x , N .VOffsetTFlags )
823- self .head = self .head - N .VOffsetTFlags .bytewidth
824- encode .Write (packer .voffset , self .Bytes , self .Head (), x )
862+ new_head = self .head - N .VOffsetTFlags .bytewidth
863+ self .head = UOffsetTFlags .py_type (new_head )
864+ encode .Write (packer .voffset , self .Bytes , new_head , x )
825865
826866 def PlaceSOffsetT (self , x ):
827867 """PlaceSOffsetT prepends a SOffsetT to the Builder, without checking
828868
829869 for space.
830870 """
831871 N .enforce_number (x , N .SOffsetTFlags )
832- self .head = self .head - N .SOffsetTFlags .bytewidth
833- encode .Write (packer .soffset , self .Bytes , self .Head (), x )
872+ new_head = self .head - N .SOffsetTFlags .bytewidth
873+ self .head = UOffsetTFlags .py_type (new_head )
874+ encode .Write (packer .soffset , self .Bytes , new_head , x )
834875
835876 def PlaceUOffsetT (self , x ):
836877 """PlaceUOffsetT prepends a UOffsetT to the Builder, without checking
837878
838879 for space.
839880 """
840881 N .enforce_number (x , N .UOffsetTFlags )
841- self .head = self .head - N .UOffsetTFlags .bytewidth
842- encode .Write (packer .uoffset , self .Bytes , self .Head (), x )
882+ new_head = self .head - N .UOffsetTFlags .bytewidth
883+ self .head = UOffsetTFlags .py_type (new_head )
884+ encode .Write (packer .uoffset , self .Bytes , new_head , x )
843885
844886 ## @endcond
845887
0 commit comments