Skip to content

Commit 6b89193

Browse files
methanemethane
committed
fallback: Use struct.unpack_from when possible
1 parent b78c0c5 commit 6b89193

File tree

1 file changed

+100
-50
lines changed

1 file changed

+100
-50
lines changed

‎msgpack/fallback.py‎

Lines changed: 100 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -265,14 +265,18 @@ def read_bytes(self, n):
265265
return self._read(n)
266266

267267
def _read(self, n):
268-
# (int, Optional[Callable]) -> bytearray
268+
# (int) -> bytearray
269+
self._reserve(n)
270+
i = self._buff_i
271+
self._buff_i = i+n
272+
return self._buffer[i:i+n]
273+
274+
def _reserve(self, n):
269275
remain_bytes = len(self._buffer) - self._buff_i - n
270276

271277
# Fast path: buffer has n bytes already
272278
if remain_bytes >= 0:
273-
ret = self._buffer[self._buff_i:self._buff_i+n]
274-
self._buff_i += n
275-
return ret
279+
return
276280

277281
if self._feeding:
278282
self._buff_i = self._buf_checkpoint
@@ -299,33 +303,23 @@ def _read(self, n):
299303
self._buff_i = 0 # rollback
300304
raise OutOfData
301305

302-
if len(self._buffer) == n:
303-
# checkpoint == 0
304-
ret = self._buffer
305-
self._buffer = b""
306-
self._buff_i = 0
307-
else:
308-
ret = self._buffer[self._buff_i:self._buff_i+n]
309-
self._buff_i += n
310-
311-
return ret
312-
313306
def _read_header(self, execute=EX_CONSTRUCT):
314307
typ = TYPE_IMMEDIATE
315308
n = 0
316309
obj = None
317-
c = self._read(1)
318-
b = ord(c)
319-
if b & 0b10000000 == 0:
310+
self._reserve(1)
311+
b = struct.unpack_from("B", self._buffer, self._buff_i)[0]
312+
self._buff_i += 1
313+
if b & 0b10000000 == 0:
320314
obj = b
321315
elif b & 0b11100000 == 0b11100000:
322-
obj = struct.unpack("b", c)[0]
316+
obj = -1 - (b ^ 0xff)
323317
elif b & 0b11100000 == 0b10100000:
324318
n = b & 0b00011111
325-
obj = self._read(n)
326319
typ = TYPE_RAW
327320
if n > self._max_str_len:
328321
raise UnpackValueError("%s exceeds max_str_len(%s)", n, self._max_str_len)
322+
obj = self._read(n)
329323
elif b & 0b11110000 == 0b10010000:
330324
n = b & 0b00001111
331325
typ = TYPE_ARRAY
@@ -344,120 +338,176 @@ def _read_header(self, execute=EX_CONSTRUCT):
344338
obj = True
345339
elif b == 0xc4:
346340
typ = TYPE_BIN
347-
n = struct.unpack("B", self._read(1))[0]
341+
self._reserve(1)
342+
n = struct.unpack_from("B", self._buffer, self._buff_i)[0]
343+
self._buff_i += 1
348344
if n > self._max_bin_len:
349345
raise UnpackValueError("%s exceeds max_bin_len(%s)" % (n, self._max_bin_len))
350346
obj = self._read(n)
351347
elif b == 0xc5:
352348
typ = TYPE_BIN
353-
n = struct.unpack(">H", self._read(2))[0]
349+
self._reserve(2)
350+
n = struct.unpack_from(">H", self._buffer, self._buff_i)[0]
351+
self._buff_i += 2
354352
if n > self._max_bin_len:
355353
raise UnpackValueError("%s exceeds max_bin_len(%s)" % (n, self._max_bin_len))
356354
obj = self._read(n)
357355
elif b == 0xc6:
358356
typ = TYPE_BIN
359-
n = struct.unpack(">I", self._read(4))[0]
357+
self._reserve(4)
358+
n = struct.unpack_from(">I", self._buffer, self._buff_i)[0]
359+
self._buff_i += 4
360360
if n > self._max_bin_len:
361361
raise UnpackValueError("%s exceeds max_bin_len(%s)" % (n, self._max_bin_len))
362362
obj = self._read(n)
363363
elif b == 0xc7: # ext 8
364364
typ = TYPE_EXT
365-
L, n = struct.unpack('Bb', self._read(2))
365+
self._reserve(2)
366+
L, n = struct.unpack_from('Bb', self._buffer, self._buff_i)
367+
self._buff_i += 2
366368
if L > self._max_ext_len:
367369
raise UnpackValueError("%s exceeds max_ext_len(%s)" % (L, self._max_ext_len))
368370
obj = self._read(L)
369371
elif b == 0xc8: # ext 16
370372
typ = TYPE_EXT
371-
L, n = struct.unpack('>Hb', self._read(3))
373+
self._reserve(3)
374+
L, n = struct.unpack_from('>Hb', self._buffer, self._buff_i)
375+
self._buff_i += 3
372376
if L > self._max_ext_len:
373377
raise UnpackValueError("%s exceeds max_ext_len(%s)" % (L, self._max_ext_len))
374378
obj = self._read(L)
375379
elif b == 0xc9: # ext 32
376380
typ = TYPE_EXT
377-
L, n = struct.unpack('>Ib', self._read(5))
381+
self._reserve(5)
382+
L, n = struct.unpack_from('>Ib', self._buffer, self._buff_i)
383+
self._buff_i += 5
378384
if L > self._max_ext_len:
379385
raise UnpackValueError("%s exceeds max_ext_len(%s)" % (L, self._max_ext_len))
380386
obj = self._read(L)
381387
elif b == 0xca:
382-
obj = struct.unpack(">f", self._read(4))[0]
388+
self._reserve(4)
389+
obj = struct.unpack_from(">f", self._buffer, self._buff_i)[0]
390+
self._buff_i += 4
383391
elif b == 0xcb:
384-
obj = struct.unpack(">d", self._read(8))[0]
392+
self._reserve(8)
393+
obj = struct.unpack_from(">d", self._buffer, self._buff_i)[0]
394+
self._buff_i += 8
385395
elif b == 0xcc:
386-
obj = struct.unpack("B", self._read(1))[0]
396+
self._reserve(1)
397+
obj = struct.unpack_from("B", self._buffer, self._buff_i)[0]
398+
self._buff_i += 1
387399
elif b == 0xcd:
388-
obj = struct.unpack(">H", self._read(2))[0]
400+
self._reserve(2)
401+
obj = struct.unpack_from(">H", self._buffer, self._buff_i)[0]
402+
self._buff_i += 2
389403
elif b == 0xce:
390-
obj = struct.unpack(">I", self._read(4))[0]
404+
self._reserve(4)
405+
obj = struct.unpack_from(">I", self._buffer, self._buff_i)[0]
406+
self._buff_i += 4
391407
elif b == 0xcf:
392-
obj = struct.unpack(">Q", self._read(8))[0]
408+
self._reserve(8)
409+
obj = struct.unpack_from(">Q", self._buffer, self._buff_i)[0]
410+
self._buff_i += 8
393411
elif b == 0xd0:
394-
obj = struct.unpack("b", self._read(1))[0]
412+
self._reserve(1)
413+
obj = struct.unpack_from("b", self._buffer, self._buff_i)[0]
414+
self._buff_i += 1
395415
elif b == 0xd1:
396-
obj = struct.unpack(">h", self._read(2))[0]
416+
self._reserve(2)
417+
obj = struct.unpack_from(">h", self._buffer, self._buff_i)[0]
418+
self._buff_i += 2
397419
elif b == 0xd2:
398-
obj = struct.unpack(">i", self._read(4))[0]
420+
self._reserve(4)
421+
obj = struct.unpack_from(">i", self._buffer, self._buff_i)[0]
422+
self._buff_i += 4
399423
elif b == 0xd3:
400-
obj = struct.unpack(">q", self._read(8))[0]
424+
self._reserve(8)
425+
obj = struct.unpack_from(">q", self._buffer, self._buff_i)[0]
426+
self._buff_i += 8
401427
elif b == 0xd4: # fixext 1
402428
typ = TYPE_EXT
403429
if self._max_ext_len < 1:
404430
raise UnpackValueError("%s exceeds max_ext_len(%s)" % (1, self._max_ext_len))
405-
n, obj = struct.unpack('b1s', self._read(2))
431+
self._reserve(2)
432+
n, obj = struct.unpack_from("b1s", self._buffer, self._buff_i)
433+
self._buff_i += 2
406434
elif b == 0xd5: # fixext 2
407435
typ = TYPE_EXT
408436
if self._max_ext_len < 2:
409437
raise UnpackValueError("%s exceeds max_ext_len(%s)" % (2, self._max_ext_len))
410-
n, obj = struct.unpack('b2s', self._read(3))
438+
self._reserve(3)
439+
n, obj = struct.unpack_from("b2s", self._buffer, self._buff_i)
440+
self._buff_i += 3
411441
elif b == 0xd6: # fixext 4
412442
typ = TYPE_EXT
413443
if self._max_ext_len < 4:
414444
raise UnpackValueError("%s exceeds max_ext_len(%s)" % (4, self._max_ext_len))
415-
n, obj = struct.unpack('b4s', self._read(5))
445+
self._reserve(5)
446+
n, obj = struct.unpack_from("b4s", self._buffer, self._buff_i)
447+
self._buff_i += 5
416448
elif b == 0xd7: # fixext 8
417449
typ = TYPE_EXT
418450
if self._max_ext_len < 8:
419451
raise UnpackValueError("%s exceeds max_ext_len(%s)" % (8, self._max_ext_len))
420-
n, obj = struct.unpack('b8s', self._read(9))
452+
self._reserve(9)
453+
n, obj = struct.unpack_from("b8s", self._buffer, self._buff_i)
454+
self._buff_i += 9
421455
elif b == 0xd8: # fixext 16
422456
typ = TYPE_EXT
423457
if self._max_ext_len < 16:
424458
raise UnpackValueError("%s exceeds max_ext_len(%s)" % (16, self._max_ext_len))
425-
n, obj = struct.unpack('b16s', self._read(17))
459+
self._reserve(17)
460+
n, obj = struct.unpack_from("b16s", self._buffer, self._buff_i)
461+
self._buff_i += 17
426462
elif b == 0xd9:
427463
typ = TYPE_RAW
428-
n = struct.unpack("B", self._read(1))[0]
464+
self._reserve(1)
465+
n, = struct.unpack_from("B", self._buffer, self._buff_i)
466+
self._buff_i += 1
429467
if n > self._max_str_len:
430468
raise UnpackValueError("%s exceeds max_str_len(%s)", n, self._max_str_len)
431469
obj = self._read(n)
432470
elif b == 0xda:
433471
typ = TYPE_RAW
434-
n = struct.unpack(">H", self._read(2))[0]
472+
self._reserve(2)
473+
n, = struct.unpack_from(">H", self._buffer, self._buff_i)
474+
self._buff_i += 2
435475
if n > self._max_str_len:
436476
raise UnpackValueError("%s exceeds max_str_len(%s)", n, self._max_str_len)
437477
obj = self._read(n)
438478
elif b == 0xdb:
439479
typ = TYPE_RAW
440-
n = struct.unpack(">I", self._read(4))[0]
480+
self._reserve(4)
481+
n, = struct.unpack_from(">I", self._buffer, self._buff_i)
482+
self._buff_i += 4
441483
if n > self._max_str_len:
442484
raise UnpackValueError("%s exceeds max_str_len(%s)", n, self._max_str_len)
443485
obj = self._read(n)
444486
elif b == 0xdc:
445-
n = struct.unpack(">H", self._read(2))[0]
487+
typ = TYPE_ARRAY
488+
self._reserve(2)
489+
n, = struct.unpack_from(">H", self._buffer, self._buff_i)
490+
self._buff_i += 2
446491
if n > self._max_array_len:
447492
raise UnpackValueError("%s exceeds max_array_len(%s)", n, self._max_array_len)
448-
typ = TYPE_ARRAY
449493
elif b == 0xdd:
450-
n = struct.unpack(">I", self._read(4))[0]
494+
typ = TYPE_ARRAY
495+
self._reserve(4)
496+
n, = struct.unpack_from(">I", self._buffer, self._buff_i)
497+
self._buff_i += 4
451498
if n > self._max_array_len:
452499
raise UnpackValueError("%s exceeds max_array_len(%s)", n, self._max_array_len)
453-
typ = TYPE_ARRAY
454500
elif b == 0xde:
455-
n = struct.unpack(">H", self._read(2))[0]
501+
self._reserve(2)
502+
n, = struct.unpack_from(">H", self._buffer, self._buff_i)
503+
self._buff_i += 2
456504
if n > self._max_map_len:
457505
raise UnpackValueError("%s exceeds max_map_len(%s)", n, self._max_map_len)
458506
typ = TYPE_MAP
459507
elif b == 0xdf:
460-
n = struct.unpack(">I", self._read(4))[0]
508+
self._reserve(4)
509+
n, = struct.unpack_from(">I", self._buffer, self._buff_i)
510+
self._buff_i += 4
461511
if n > self._max_map_len:
462512
raise UnpackValueError("%s exceeds max_map_len(%s)", n, self._max_map_len)
463513
typ = TYPE_MAP

0 commit comments

Comments
 (0)