8000 Merge pull request #199 from methane/struct-unpack-from · sesam-io/msgpack-python@b887c1a · GitHub
[go: up one dir, main page]
Skip to content

Commit b887c1a

Browse files
methanemethane
committed
Merge pull request msgpack#199 from methane/struct-unpack-from
Use struct.unpack_from instead of struct.unpack
2 parents b78c0c5 + c16a1c6 commit b887c1a

File tree

1 file changed

+111
-55
lines changed

1 file changed

+111
-55
lines changed

‎msgpack/fallback.py

Lines changed: 111 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -196,7 +196,7 @@ def __init__(self, file_like=None, read_size=0, use_list=True,
196196
self._feeding = False
197197

198198
#: array of bytes feeded.
199-
self._buffer = b""
199+
self._buffer = bytearray()
200200
#: Which position we currently reads
201201
self._buff_i = 0
202202

@@ -249,7 +249,7 @@ def feed(self, next_bytes):
249249
raise BufferFull
250250
# bytes + bytearray -> bytearray
251251
# So cast before append
252-
self._buffer += bytes(next_bytes)
252+
self._buffer += next_bytes
253253

254254
def _consume(self):
255255
""" Gets rid of the used parts of the buffer. """
@@ -265,22 +265,26 @@ 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
279283
raise OutOfData
280284

281285
# Strip buffer before checkpoint before reading file.
282286
if self._buf_checkpoint > 0:
283-
self._buffer = self._buffer[self._buf_checkpoint:]
287+
del self._buffer[:self._buf_checkpoint]
284288
self._buff_i -= self._buf_checkpoint
285289
self._buf_checkpoint = 0
286290

@@ -299,33 +303,24 @@ 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+
b = self._buffer[self._buff_i]
313+
self._buff_i += 1
314+
if b & 0b10000000 == 0:
320315
obj = b
321316
elif b & 0b11100000 == 0b11100000:
322-
obj = struct.unpack("b", c)[0]
317+
obj = -1 - (b ^ 0xff)
323318
elif b & 0b11100000 == 0b10100000:
324319
n = b & 0b00011111
325-
obj = self._read(n)
326320
typ = TYPE_RAW
327321
if n > self._max_str_len:
328322
raise UnpackValueError("%s exceeds max_str_len(%s)", n, self._max_str_len)
323+
obj = self._read(n)
329324
elif b & 0b11110000 == 0b10010000:
330325
n = b & 0b00001111
331326
typ = TYPE_ARRAY
@@ -344,120 +339,179 @@ def _read_header(self, execute=EX_CONSTRUCT):
344339
obj = True
345340
elif b == 0xc4:
346341
typ = TYPE_BIN
347-
n = struct.unpack("B", self._read(1))[0]
342+
self._reserve(1)
343+
#n = struct.unpack_from("B", self._buffer, self._buff_i)[0]
344+
n = self._buffer[self._buff_i]
345+
self._buff_i += 1
348346
if n > self._max_bin_len:
349347
raise UnpackValueError("%s exceeds max_bin_len(%s)" % (n, self._max_bin_len))
350348
obj = self._read(n)
351349
elif b == 0xc5:
352350
typ = TYPE_BIN
353-
n = struct.unpack(">H", self._read(2))[0]
351+
self._reserve(2)
352+
n = struct.unpack_from(">H", self._buffer, self._buff_i)[0]
353+
self._buff_i += 2
354354
if n > self._max_bin_len:
355355
raise UnpackValueError("%s exceeds max_bin_len(%s)" % (n, self._max_bin_len))
356356
obj = self._read(n)
357357
elif b == 0xc6:
358358
typ = TYPE_BIN
359-
n = struct.unpack(">I", self._read(4))[0]
359+
self._reserve(4)
360+
n = struct.unpack_from(">I", self._buffer, self._buff_i)[0]
361+
self._buff_i += 4
360362
if n > self._max_bin_len:
361363
raise UnpackValueError("%s exceeds max_bin_len(%s)" % (n, self._max_bin_len))
362364
obj = self._read(n)
363365
elif b == 0xc7: # ext 8
364366
typ = TYPE_EXT
365-
L, n = struct.unpack('Bb', self._read(2))
367+
self._reserve(2)
368+
L, n = struct.unpack_from('Bb', self._buffer, self._buff_i)
369+
self._buff_i += 2
366370
if L > self._max_ext_len:
367371
raise UnpackValueError("%s exceeds max_ext_len(%s)" % (L, self._max_ext_len))
368372
obj = self._read(L)
369373
elif b == 0xc8: # ext 16
370374
typ = TYPE_EXT
371-
L, n = struct.unpack('>Hb', self._read(3))
375+
self._reserve(3)
376+
L, n = struct.unpack_from('>Hb', self._buffer, self._buff_i)
377+
self._buff_i += 3
372378
if L > self._max_ext_len:
373379
raise UnpackValueError("%s exceeds max_ext_len(%s)" % (L, self._max_ext_len))
374380
obj = self._read(L)
375381
elif b == 0xc9: # ext 32
376382
typ = TYPE_EXT
377-
L, n = struct.unpack('>Ib', self._read(5))
383+
self._reserve(5)
384+
L, n = struct.unpack_from('>Ib', self._buffer, self._buff_i)
385+
self._buff_i += 5
378386
if L > self._max_ext_len:
379387
raise UnpackValueError("%s exceeds max_ext_len(%s)" % (L, self._max_ext_len))
380388
obj = self._read(L)
381389
elif b == 0xca:
382-
obj = struct.unpack(">f", self._read(4))[0]
390+
self._reserve(4)
391+
obj = struct.unpack_from(">f", self._buffer, self._buff_i)[0]
392+
self._buff_i += 4
383393
elif b == 0xcb:
384-
obj = struct.unpack(">d", self._read(8))[0]
394+
self._reserve(8)
395+
obj = struct.unpack_from(">d", self._buffer, self._buff_i)[0]
396+
self._buff_i += 8
385397
elif b == 0xcc:
386-
obj = struct.unpack("B", self._read(1))[0]
398+
self._reserve(1)
399+
#obj = struct.unpack_from("B", self._buffer, self._buff_i)[0]
400+
obj = self._buffer[self._buff_i]
401+
self._buff_i += 1
387402
elif b == 0xcd:
388-
obj = struct.unpack(">H", self._read(2))[0]
403+
self._reserve(2)
404+
obj = struct.unpack_from(">H", self._buffer, self._buff_i)[0]
405+
self._buff_i += 2
389406
elif b == 0xce:
390-
obj = struct.unpack(">I", self._read(4))[0]
407+
self._reserve(4)
408+
obj = struct.unpack_from(">I", self._buffer, self._buff_i)[0]
409+
self._buff_i += 4
391410
elif b == 0xcf:
392-
obj = struct.unpack(">Q", self._read(8))[0]
411+
self._reserve(8)
412+
obj = struct.unpack_from(">Q", self._buffer, self._buff_i)[0]
413+
self._buff_i += 8
393414
elif b == 0xd0:
394-
obj = struct.unpack("b", self._read(1))[0]
415+
self._reserve(1)
416+
obj = struct.unpack_from("b", self._buffer, self._buff_i)[0]
417+
self._buff_i += 1
395418
elif b == 0xd1:
396-
obj = struct.unpack(">h", self._read(2))[0]
419+
self._reserve(2)
420+
obj = struct.unpack_from(">h", self._buffer, self._buff_i)[0]
421+
self._buff_i += 2
397422
elif b == 0xd2:
398-
obj = struct.unpack(">i", self._read(4))[0]
423+
self._reserve(4)
424+
obj = struct.unpack_from(">i", self._buffer, self._buff_i)[0]
425+
self._buff_i += 4
399426
elif b == 0xd3:
400-
obj = struct.unpack(">q", self._read(8))[0]
427+
self._reserve(8)
428+
obj = struct.unpack_from(">q", self._buffer, self._buff_i)[0]
429+
self._buff_i += 8
401430
elif b == 0xd4: # fixext 1
402431
typ = TYPE_EXT
403432
if self._max_ext_len < 1:
404433
raise UnpackValueError("%s exceeds max_ext_len(%s)" % (1, self._max_ext_len))
405-
n, obj = struct.unpack('b1s', self._read(2))
434+
self._reserve(2)
435+
n, obj = struct.unpack_from("b1s", self._buffer, self._buff_i)
436+
self._buff_i += 2
406437
elif b == 0xd5: # fixext 2
407438
typ = TYPE_EXT
408439
if self._max_ext_len < 2:
409440
raise UnpackValueError("%s exceeds max_ext_len(%s)" % (2, self._max_ext_len))
410-
n, obj = struct.unpack('b2s', self._read(3))
441+
self._reserve(3)
442+
n, obj = struct.unpack_from("b2s", self._buffer, self._buff_i)
443+
self._buff_i += 3
411444
elif b == 0xd6: # fixext 4
412445
typ = TYPE_EXT
413446
if self._max_ext_len < 4:
414447
raise UnpackValueError("%s exceeds max_ext_len(%s)" % (4, self._max_ext_len))
415-
n, obj = struct.unpack('b4s', self._read(5))
448+
self._reserve(5)
449+
n, obj = struct.unpack_from("b4s", self._buffer, self._buff_i)
450+
self._buff_i += 5
416451
elif b == 0xd7: # fixext 8
417452
typ = TYPE_EXT
418453
if self._max_ext_len < 8:
419454
raise UnpackValueError("%s exceeds max_ext_len(%s)" % (8, self._max_ext_len))
420-
n, obj = struct.unpack('b8s', self._read(9))
455+
self._reserve(9)
456+
n, obj = struct.unpack_from("b8s", self._buffer, self._buff_i)
457+
self._buff_i += 9
421458
elif b == 0xd8: # fixext 16
422459
typ = TYPE_EXT
423460
if self._max_ext_len < 16:
424461
raise UnpackValueError("%s exceeds max_ext_len(%s)" % (16, self._max_ext_len))
425-
n, obj = struct.unpack('b16s', self._read(17))
462+
self._reserve(17)
463+
n, obj = struct.unpack_from("b16s", self._buffer, self._buff_i)
464+
self._buff_i += 17
426465
elif b == 0xd9:
427466
typ = TYPE_RAW
428-
n = struct.unpack("B", self._read(1))[0]
467+
self._reserve(1)
468+
#n, = struct.unpack_from("B", self._buffer, self._buff_i)
469+
n = self._buffer[self._buff_i]
470+
self._buff_i += 1
429471
if n > self._max_str_len:
430472
raise UnpackValueError("%s exceeds max_str_len(%s)", n, self._max_str_len)
431473
obj = self._read(n)
432474
elif b == 0xda:
433475
typ = TYPE_RAW
434-
n = struct.unpack(">H", self._read(2))[0]
476+
self._reserve(2)
477+
n, = struct.unpack_from(">H", self._buffer, self._buff_i)
478+
self._buff_i += 2
435479
if n > self._max_str_len:
436480
raise UnpackValueError("%s exceeds max_str_len(%s)", n, self._max_str_len)
437481
obj = self._read(n)
438482
elif b == 0xdb:
439483
typ = TYPE_RAW
440-
n = struct.unpack(">I", self._read(4))[0]
484+
self._reserve(4)
485+
n, = struct.unpack_from(">I", self._buffer, self._buff_i)
486+
self._buff_i += 4
441487
if n > self._max_str_len:
442488
raise UnpackValueError("%s exceeds max_str_len(%s)", n, self._max_str_len)
443489
obj = self._read(n)
444490
elif b == 0xdc:
445-
n = struct.unpack(">H", self._read(2))[0]
491+
typ = TYPE_ARRAY
492+
self._reserve(2)
493+
n, = struct.unpack_from(">H", self._buffer, self._buff_i)
494+
self._buff_i += 2
446495
if n > self._max_array_len:
447496
raise UnpackValueError("%s exceeds max_array_len(%s)", n, self._max_array_len)
448-
typ = TYPE_ARRAY
449497
elif b == 0xdd:
450-
n = struct.unpack(">I", self._read(4))[0]
498+
typ = TYPE_ARRAY
499+
self._reserve(4)
500+
n, = struct.unpack_from(">I", self._buffer, self._buff_i)
501+
self._buff_i += 4
451502
if n > self._max_array_len:
452503
raise UnpackValueError("%s exceeds max_array_len(%s)", n, self._max_array_len)
453-
typ = TYPE_ARRAY
454504
elif b == 0xde:
455-
n = struct.unpack(">H", self._read(2))[0]
505+
self._reserve(2)
506+
n, = struct.unpack_from(">H", self._buffer, self._buff_i)
507+
self._buff_i += 2
456508
if n > self._max_map_len:
457509
raise UnpackValueError("%s exceeds max_map_len(%s)", n, self._max_map_len)
458510
typ = TYPE_MAP
459511
elif b == 0xdf:
460-
n = struct.unpack(">I", self._read(4))[0]
512+
self._reserve(4)
513+
n, = struct.unpack_from(">I", self._buffer, self._buff_i)
514+
self._buff_i += 4
461515
if n > self._max_map_len:
462516
raise UnpackValueError("%s exceeds max_map_len(%s)", n, self._max_map_len)
463517
typ = TYPE_MAP
@@ -515,11 +569,13 @@ def _unpack(self, execute=EX_CONSTRUCT):
515569
if typ == TYPE_RAW:
516570
if self._encoding is not None:
517571
obj = obj.decode(self._encoding, self._unicode_errors)
572+
else:
573+
obj = bytes(obj)
518574
return obj
519575
if typ == TYPE_EXT:
520-
return self._ext_hook(n, obj)
576+
return self._ext_hook(n, bytes(obj))
521577
if typ == TYPE_BIN:
522-
return obj
578+
return bytes(obj)
523579
assert typ == TYPE_IMMEDIATE
524580
return obj
525581

0 commit comments

Comments
 (0)
0