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 ()
@@ -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:
@@ -313,6 +315,18 @@ def WriteVtable(self):
313315 self .current_vtable = None
314316 return objectOffset
315317
318+ ## @cond FLATBUFFERS_INTERNAL
319+ def writeVtableEntries (self , entries ):
320+ """Write a contiguous block of VOffsetT values with a single prep call."""
321+ count = len (entries )
322+ if count == 0 :
323+ return
324+ elem_size = N .VOffsetTFlags .bytewidth
325+ total_bytes = elem_size * count
326+ self .Prep (N .VOffsetTFlags .bytewidth , total_bytes - elem_size )
327+ for value in entries :
328+ self .PlaceVOffsetT (value )
329+
316330 def EndObject (self ):
317331 """EndObject writes data necessary to finish object construction."""
318332 self .assertNested ()
@@ -350,12 +364,15 @@ def Head(self):
350364 ## @cond FLATBUFFERS_INTERNAL
351365 def Offset (self ):
352366 """Offset relative to the end of the buffer."""
353- return UOffsetTFlags .py_type (len (self .Bytes ) - self .Head () )
367+ return UOffsetTFlags .py_type (len (self .Bytes ) - self .head )
354368
355369 def Pad (self , n ):
356370 """Pad places zeros at the current offset."""
357- for i in range_func (n ):
358- self .Place (0 , N .Uint8Flags )
371+ if n <= 0 :
372+ return
373+ new_head = self .head - n
374+ self .Bytes [new_head : self .head ] = b"\x00 " * n
375+ self .head = UOffsetTFlags .py_type (new_head )
359376
360377 def Prep (self , size , additionalBytes ):
361378 """Prep prepares to write an element of `size` after `additional_bytes`
@@ -372,15 +389,19 @@ def Prep(self, size, additionalBytes):
372389
373390 # Find the amount of alignment needed such that `size` is properly
374391 # aligned after `additionalBytes`:
375- alignSize = (~ (len (self .Bytes ) - self .Head () + additionalBytes )) + 1
392+ head = self .head
393+ buf_len = len (self .Bytes )
394+ alignSize = (~ (buf_len - head + additionalBytes )) + 1
376395 alignSize &= size - 1
377396
378397 # Reallocate the buffer if needed:
379- while self .Head () < alignSize + size + additionalBytes :
380- oldBufSize = len (self .Bytes )
398+ needed = alignSize + size + additionalBytes
399+ while head < needed :
400+ oldBufSize = buf_len
381401 self .growByteBuffer ()
382- updated_head = self .head + len (self .Bytes ) - oldBufSize
383- self .head = UOffsetTFlags .py_type (updated_head )
402+ buf_len = len (self .Bytes )
403+ head += buf_len - oldBufSize
404+ self .head = UOffsetTFlags .py_type (head )
384405 self .Pad (alignSize )
385406
386407 def PrependSOffsetTRelative (self , off ):
@@ -478,16 +499,17 @@ def CreateString(self, s, encoding="utf-8", errors="strict"):
478499 else :
479500 raise TypeError ("non-string passed to CreateString" )
480501
481- self .Prep (N .UOffsetTFlags .bytewidth , (len (x ) + 1 ) * N .Uint8Flags .bytewidth )
502+ payload_len = len (x )
503+ self .Prep (
504+ N .UOffsetTFlags .bytewidth , (payload_len + 1 ) * N .Uint8Flags .bytewidth
505+ )
482506 self .Place (0 , N .Uint8Flags )
483507
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
508+ new_head = self .head - payload_len
509+ self .head = UOffsetTFlags .py_type (new_head )
510+ self .Bytes [new_head : new_head + payload_len ] = x
489511
490- self .vectorNumElems = len ( x )
512+ self .vectorNumElems = payload_len
491513 return self .EndVector ()
492514
493515 def CreateByteVector (self , x ):
@@ -501,15 +523,13 @@ def CreateByteVector(self, x):
501523 if not isinstance (x , compat .binary_types ):
502524 raise TypeError ("non-byte vector passed to CreateByteVector" )
503525
504- self .Prep (N .UOffsetTFlags .bytewidth , len (x ) * N .Uint8Flags .bytewidth )
526+ data_len = len (x )
527+ self .Prep (N .UOffsetTFlags .bytewidth , data_len * N .Uint8Flags .bytewidth )
528+ new_head = self .head - data_len
529+ self .head = UOffsetTFlags .py_type (new_head )
530+ self .Bytes [new_head : new_head + data_len ] = x
505531
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
511-
512- self .vectorNumElems = len (x )
532+ self .vectorNumElems = data_len
513533 return self .EndVector ()
514534
515535 def CreateNumpyVector (self , x ):
@@ -537,13 +557,11 @@ def CreateNumpyVector(self, x):
537557 x_lend = x .byteswap (inplace = False )
538558
539559 # 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" )
560+ payload = x_lend .tobytes (order = "C" )
561+ payload_len = len (payload )
562+ new_head = self .head - payload_len
563+ self .head = UOffsetTFlags .py_type (new_head )
564+ self .Bytes [new_head : new_head + payload_len ] = payload
547565
548566 self .vectorNumElems = x .size
549567 return self .EndVector ()
@@ -632,8 +650,31 @@ def FinishSizePrefixed(self, rootTable, file_identifier=None):
632650
633651 ## @cond FLATBUFFERS_INTERNAL
634652 def Prepend (self , flags , off ):
635- self .Prep (flags .bytewidth , 0 )
636- self .Place (off , flags )
653+ size = flags .bytewidth
654+ if size > self .minalign :
655+ self .minalign = size
656+
657+ head = self .head
658+ buf_len = len (self .Bytes )
659+ alignSize = (~ (buf_len - head )) + 1
660+ alignSize &= size - 1
661+
662+ needed = alignSize + size
663+ while head < needed :
664+ oldBufSize = buf_len
665+ self .growByteBuffer ()
666+ buf_len = len (self .Bytes )
667+ head += buf_len - oldBufSize
668+
669+ if alignSize :
670+ new_head = head - alignSize
671+ self .Bytes [new_head :head ] = b"\x00 " * alignSize
672+ head = new_head
673+
674+ N .enforce_number (off , flags )
675+ head -= size
676+ self .head = UOffsetTFlags .py_type (head )
677+ encode .Write (flags .packer_type , self .Bytes , head , off )
637678
638679 def PrependSlot (self , flags , o , x , d ):
639680 if x is not None :
@@ -800,46 +841,46 @@ def ForceDefaults(self, forceDefaults):
800841
801842 ##############################################################
802843
803- ## @cond FLATBUFFERS_INTERNAL
804- def PrependVOffsetT (self , x ):
805- self .Prepend (N .VOffsetTFlags , x )
806-
807844 def Place (self , x , flags ):
808845 """Place prepends a value specified by `flags` to the Builder,
809846
810847 without checking for available space.
811848 """
812849
813850 N .enforce_number (x , flags )
814- self .head = self .head - flags .bytewidth
815- encode .Write (flags .packer_type , self .Bytes , self .Head (), x )
851+ new_head = self .head - flags .bytewidth
852+ self .head = UOffsetTFlags .py_type (new_head )
853+ encode .Write (flags .packer_type , self .Bytes , new_head , x )
816854
817855 def PlaceVOffsetT (self , x ):
818856 """PlaceVOffsetT prepends a VOffsetT to the Builder, without checking
819857
820858 for space.
821859 """
822860 N .enforce_number (x , N .VOffsetTFlags )
823- self .head = self .head - N .VOffsetTFlags .bytewidth
824- encode .Write (packer .voffset , self .Bytes , self .Head (), x )
861+ new_head = self .head - N .VOffsetTFlags .bytewidth
862+ self .head = UOffsetTFlags .py_type (new_head )
863+ encode .Write (packer .voffset , self .Bytes , new_head , x )
825864
826865 def PlaceSOffsetT (self , x ):
827866 """PlaceSOffsetT prepends a SOffsetT to the Builder, without checking
828867
829868 for space.
830869 """
831870 N .enforce_number (x , N .SOffsetTFlags )
832- self .head = self .head - N .SOffsetTFlags .bytewidth
833- encode .Write (packer .soffset , self .Bytes , self .Head (), x )
871+ new_head = self .head - N .SOffsetTFlags .bytewidth
872+ self .head = UOffsetTFlags .py_type (new_head )
873+ encode .Write (packer .soffset , self .Bytes , new_head , x )
834874
835875 def PlaceUOffsetT (self , x ):
836876 """PlaceUOffsetT prepends a UOffsetT to the Builder, without checking
837877
838878 for space.
839879 """
840880 N .enforce_number (x , N .UOffsetTFlags )
841- self .head = self .head - N .UOffsetTFlags .bytewidth
842- encode .Write (packer .uoffset , self .Bytes , self .Head (), x )
881+ new_head = self .head - N .UOffsetTFlags .bytewidth
882+ self .head = UOffsetTFlags .py_type (new_head )
883+ encode .Write (packer .uoffset , self .Bytes , new_head , x )
843884
844885 ## @endcond
845886
0 commit comments