Skip to content

Commit e9c42fa

Browse files
methanemethane
committed
fallback: simplify write_bytes callback implementation
1 parent 3322a76 commit e9c42fa

File tree

1 file changed

+64
-60
lines changed

1 file changed

+64
-60
lines changed

‎msgpack/fallback.py‎

Lines changed: 64 additions & 60 deletions
Original file line numberDiff line numberDiff line change
@@ -264,16 +264,14 @@ def _get_extradata(self):
264264
def read_bytes(self, n):
265265
return self._read(n)
266266

267-
def _read(self, n, write_bytes=None):
267+
def _read(self, n):
268268
# (int, Optional[Callable]) -> bytearray
269269
remain_bytes = len(self._buffer) - self._buff_i - n
270270

271271
# Fast path: buffer has n bytes already
272272
if remain_bytes >= 0:
273273
ret = self._buffer[self._buff_i:self._buff_i+n]
274274
self._buff_i += n
275-
if write_bytes is not None:
276-
write_bytes(ret)
277275
return ret
278276

279277
if self._feeding:
@@ -310,23 +308,21 @@ def _read(self, n, write_bytes=None):
310308
ret = self._buffer[self._buff_i:self._buff_i+n]
311309
self._buff_i += n
312310

313-
if write_bytes is not None:
314-
write_bytes(ret)
315311
return ret
316312

317-
def _read_header(self, execute=EX_CONSTRUCT, write_bytes=None):
313+
def _read_header(self, execute=EX_CONSTRUCT):
318314
typ = TYPE_IMMEDIATE
319315
n = 0
320316
obj = None
321-
c = self._read(1, write_bytes)
317+
c = self._read(1)
322318
b = ord(c)
323319
if b & 0b10000000 == 0:
324320
obj = b
325321
elif b & 0b11100000 == 0b11100000:
326322
obj = struct.unpack("b", c)[0]
327323
elif b & 0b11100000 == 0b10100000:
328324
n = b & 0b00011111
329-
obj = self._read(n, write_bytes)
325+
obj = self._read(n)
330326
typ = TYPE_RAW
331327
if n > self._max_str_len:
332328
raise UnpackValueError("%s exceeds max_str_len(%s)", n, self._max_str_len)
@@ -348,129 +344,129 @@ def _read_header(self, execute=EX_CONSTRUCT, write_bytes=None):
348344
obj = True
349345
elif b == 0xc4:
350346
typ = TYPE_BIN
351-
n = struct.unpack("B", self._read(1, write_bytes))[0]
347+
n = struct.unpack("B", self._read(1))[0]
352348
if n > self._max_bin_len:
353349
raise UnpackValueError("%s exceeds max_bin_len(%s)" % (n, self._max_bin_len))
354-
obj = self._read(n, write_bytes)
350+
obj = self._read(n)
355351
elif b == 0xc5:
356352
typ = TYPE_BIN
357-
n = struct.unpack(">H", self._read(2, write_bytes))[0]
353+
n = struct.unpack(">H", self._read(2))[0]
358354
if n > self._max_bin_len:
359355
raise UnpackValueError("%s exceeds max_bin_len(%s)" % (n, self._max_bin_len))
360-
obj = self._read(n, write_bytes)
356+
obj = self._read(n)
361357
elif b == 0xc6:
362358
typ = TYPE_BIN
363-
n = struct.unpack(">I", self._read(4, write_bytes))[0]
359+
n = struct.unpack(">I", self._read(4))[0]
364360
if n > self._max_bin_len:
365361
raise UnpackValueError("%s exceeds max_bin_len(%s)" % (n, self._max_bin_len))
366-
obj = self._read(n, write_bytes)
362+
obj = self._read(n)
367363
elif b == 0xc7: # ext 8
368364
typ = TYPE_EXT
369-
L, n = struct.unpack('Bb', self._read(2, write_bytes))
365+
L, n = struct.unpack('Bb', self._read(2))
370366
if L > self._max_ext_len:
371367
raise UnpackValueError("%s exceeds max_ext_len(%s)" % (L, self._max_ext_len))
372-
obj = self._read(L, write_bytes)
368+
obj = self._read(L)
373369
elif b == 0xc8: # ext 16
374370
typ = TYPE_EXT
375-
L, n = struct.unpack('>Hb', self._read(3, write_bytes))
371+
L, n = struct.unpack('>Hb', self._read(3))
376372
if L > self._max_ext_len:
377373
raise UnpackValueError("%s exceeds max_ext_len(%s)" % (L, self._max_ext_len))
378-
obj = self._read(L, write_bytes)
374+
obj = self._read(L)
379375
elif b == 0xc9: # ext 32
380376
typ = TYPE_EXT
381-
L, n = struct.unpack('>Ib', self._read(5, write_bytes))
377+
L, n = struct.unpack('>Ib', self._read(5))
382378
if L > self._max_ext_len:
383379
raise UnpackValueError("%s exceeds max_ext_len(%s)" % (L, self._max_ext_len))
384-
obj = self._read(L, write_bytes)
380+
obj = self._read(L)
385381
elif b == 0xca:
386-
obj = struct.unpack(">f", self._read(4, write_bytes))[0]
382+
obj = struct.unpack(">f", self._read(4))[0]
387383
elif b == 0xcb:
388-
obj = struct.unpack(">d", self._read(8, write_bytes))[0]
384+
obj = struct.unpack(">d", self._read(8))[0]
389385
elif b == 0xcc:
390-
obj = struct.unpack("B", self._read(1, write_bytes))[0]
386+
obj = struct.unpack("B", self._read(1))[0]
391387
elif b == 0xcd:
392-
obj = struct.unpack(">H", self._read(2, write_bytes))[0]
388+
obj = struct.unpack(">H", self._read(2))[0]
393389
elif b == 0xce:
394-
obj = struct.unpack(">I", self._read(4, write_bytes))[0]
390+
obj = struct.unpack(">I", self._read(4))[0]
395391
elif b == 0xcf:
396-
obj = struct.unpack(">Q", self._read(8, write_bytes))[0]
392+
obj = struct.unpack(">Q", self._read(8))[0]
397393
elif b == 0xd0:
398-
obj = struct.unpack("b", self._read(1, write_bytes))[0]
394+
obj = struct.unpack("b", self._read(1))[0]
399395
elif b == 0xd1:
400-
obj = struct.unpack(">h", self._read(2, write_bytes))[0]
396+
obj = struct.unpack(">h", self._read(2))[0]
401397
elif b == 0xd2:
402-
obj = struct.unpack(">i", self._read(4, write_bytes))[0]
398+
obj = struct.unpack(">i", self._read(4))[0]
403399
elif b == 0xd3:
404-
obj = struct.unpack(">q", self._read(8, write_bytes))[0]
400+
obj = struct.unpack(">q", self._read(8))[0]
405401
elif b == 0xd4: # fixext 1
406402
typ = TYPE_EXT
407403
if self._max_ext_len < 1:
408404
raise UnpackValueError("%s exceeds max_ext_len(%s)" % (1, self._max_ext_len))
409-
n, obj = struct.unpack('b1s', self._read(2, write_bytes))
405+
n, obj = struct.unpack('b1s', self._read(2))
410406
elif b == 0xd5: # fixext 2
411407
typ = TYPE_EXT
412408
if self._max_ext_len < 2:
413409
raise UnpackValueError("%s exceeds max_ext_len(%s)" % (2, self._max_ext_len))
414-
n, obj = struct.unpack('b2s', self._read(3, write_bytes))
410+
n, obj = struct.unpack('b2s', self._read(3))
415411
elif b == 0xd6: # fixext 4
416412
typ = TYPE_EXT
417413
if self._max_ext_len < 4:
418414
raise UnpackValueError("%s exceeds max_ext_len(%s)" % (4, self._max_ext_len))
419-
n, obj = struct.unpack('b4s', self._read(5, write_bytes))
415+
n, obj = struct.unpack('b4s', self._read(5))
420416
elif b == 0xd7: # fixext 8
421417
typ = TYPE_EXT
422418
if self._max_ext_len < 8:
423419
raise UnpackValueError("%s exceeds max_ext_len(%s)" % (8, self._max_ext_len))
424-
n, obj = struct.unpack('b8s', self._read(9, write_bytes))
420+
n, obj = struct.unpack('b8s', self._read(9))
425421
elif b == 0xd8: # fixext 16
426422
typ = TYPE_EXT
427423
if self._max_ext_len < 16:
428424
raise UnpackValueError("%s exceeds max_ext_len(%s)" % (16, self._max_ext_len))
429-
n, obj = struct.unpack('b16s', self._read(17, write_bytes))
425+
n, obj = struct.unpack('b16s', self._read(17))
430426
elif b == 0xd9:
431427
typ = TYPE_RAW
432-
n = struct.unpack("B", self._read(1, write_bytes))[0]
428+
n = struct.unpack("B", self._read(1))[0]
433429
if n > self._max_str_len:
434430
raise UnpackValueError("%s exceeds max_str_len(%s)", n, self._max_str_len)
435-
obj = self._read(n, write_bytes)
431+
obj = self._read(n)
436432
elif b == 0xda:
437433
typ = TYPE_RAW
438-
n = struct.unpack(">H", self._read(2, write_bytes))[0]
434+
n = struct.unpack(">H", self._read(2))[0]
439435
if n > self._max_str_len:
440436
raise UnpackValueError("%s exceeds max_str_len(%s)", n, self._max_str_len)
441-
obj = self._read(n, write_bytes)
437+
obj = self._read(n)
442438
elif b == 0xdb:
443439
typ = TYPE_RAW
444-
n = struct.unpack(">I", self._read(4, write_bytes))[0]
440+
n = struct.unpack(">I", self._read(4))[0]
445441
if n > self._max_str_len:
446442
raise UnpackValueError("%s exceeds max_str_len(%s)", n, self._max_str_len)
447-
obj = self._read(n, write_bytes)
443+
obj = self._read(n)
448444
elif b == 0xdc:
449-
n = struct.unpack(">H", self._read(2, write_bytes))[0]
445+
n = struct.unpack(">H", self._read(2))[0]
450446
if n > self._max_array_len:
451447
raise UnpackValueError("%s exceeds max_array_len(%s)", n, self._max_array_len)
452448
typ = TYPE_ARRAY
453449
elif b == 0xdd:
454-
n = struct.unpack(">I", self._read(4, write_bytes))[0]
450+
n = struct.unpack(">I", self._read(4))[0]
455451
if n > self._max_array_len:
456452
raise UnpackValueError("%s exceeds max_array_len(%s)", n, self._max_array_len)
457453
typ = TYPE_ARRAY
458454
elif b == 0xde:
459-
n = struct.unpack(">H", self._read(2, write_bytes))[0]
455+
n = struct.unpack(">H", self._read(2))[0]
460456
if n > self._max_map_len:
461457
raise UnpackValueError("%s exceeds max_map_len(%s)", n, self._max_map_len)
462458
typ = TYPE_MAP
463459
elif b == 0xdf:
464-
n = struct.unpack(">I", self._read(4, write_bytes))[0]
460+
n = struct.unpack(">I", self._read(4))[0]
465461
if n > self._max_map_len:
466462
raise UnpackValueError("%s exceeds max_map_len(%s)", n, self._max_map_len)
467463
typ = TYPE_MAP
468464
else:
469465
raise UnpackValueError("Unknown header: 0x%x" % b)
470466
return typ, n, obj
471467

472-
def _unpack(self, execute=EX_CONSTRUCT, write_bytes=None):
473-
typ, n, obj = self._read_header(execute, write_bytes)
468+
def _unpack(self, execute=EX_CONSTRUCT):
469+
typ, n, obj = self._read_header(execute)
474470

475471
if execute == EX_READ_ARRAY_HEADER:
476472
if typ != TYPE_ARRAY:
@@ -485,11 +481,11 @@ def _unpack(self, execute=EX_CONSTRUCT, write_bytes=None):
485481
if execute == EX_SKIP:
486482
for i in xrange(n):
487483
# TODO check whether we need to call `list_hook`
488-
self._unpack(EX_SKIP, write_bytes)
484+
self._unpack(EX_SKIP)
489485
return
490486
ret = newlist_hint(n)
491487
for i in xrange(n):
492-
ret.append(self._unpack(EX_CONSTRUCT, write_bytes))
488+
ret.append(self._unpack(EX_CONSTRUCT))
493489
if self._list_hook is not None:
494490
ret = self._list_hook(ret)
495491
# TODO is the interaction between `list_hook` and `use_list` ok?
@@ -498,19 +494,19 @@ def _unpack(self, execute=EX_CONSTRUCT, write_bytes=None):
498494
if execute == EX_SKIP:
499495
for i in xrange(n):
500496
# TODO check whether we need to call hooks
501-
self._unpack(EX_SKIP, write_bytes)
502-
self._unpack(EX_SKIP, write_bytes)
497+
self._unpack(EX_SKIP)
498+
self._unpack(EX_SKIP)
503499
return
504500
if self._object_pairs_hook is not None:
505501
ret = self._object_pairs_hook(
506-
(self._unpack(EX_CONSTRUCT, write_bytes),
507-
self._unpack(EX_CONSTRUCT, write_bytes))
502+
(self._unpack(EX_CONSTRUCT),
503+
self._unpack(EX_CONSTRUCT))
508504
for _ in xrange(n))
509505
else:
510506
ret = {}
511507
for _ in xrange(n):
512-
key = self._unpack(EX_CONSTRUCT, write_bytes)
513-
ret[key] = self._unpack(EX_CONSTRUCT, write_bytes)
508+
key = self._unpack(EX_CONSTRUCT)
509+
ret[key] = self._unpack(EX_CONSTRUCT)
514510
if self._object_hook is not None:
515511
ret = self._object_hook(ret)
516512
return ret
@@ -532,7 +528,7 @@ def __iter__(self):
532528

533529
def __next__(self):
534530
try:
535-
ret = self._unpack(EX_CONSTRUCT, None)
531+
ret = self._unpack(EX_CONSTRUCT)
536532
self._consume()
537533
return ret
538534
except OutOfData:
@@ -542,21 +538,29 @@ def __next__(self):
542538
next = __next__
543539

544540
def skip(self, write_bytes=None):
545-
self._unpack(EX_SKIP, write_bytes)
541+
self._unpack(EX_SKIP)
542+
if write_bytes is not None:
543+
write_bytes(self._buffer[self._buf_checkpoint:self._buff_i])
546544
self._consume()
547545

548546
def unpack(self, write_bytes=None):
549-
ret = self._unpack(EX_CONSTRUCT, write_bytes)
547+
ret = self._unpack(EX_CONSTRUCT)
548+
if write_bytes is not None:
549+
write_bytes(self._buffer[self._buf_checkpoint:self._buff_i])
550550
self._consume()
551551
return ret
552552

553553
def read_array_header(self, write_bytes=None):
554-
ret = self._unpack(EX_READ_ARRAY_HEADER, write_bytes)
554+
ret = self._unpack(EX_READ_ARRAY_HEADER)
555+
if write_bytes is not None:
556+
write_bytes(self._buffer[self._buf_checkpoint:self._buff_i])
555557
self._consume()
556558
return ret
557559

558560
def read_map_header(self, write_bytes=None):
559-
ret = self._unpack(EX_READ_MAP_HEADER, write_bytes)
561+
ret = self._unpack(EX_READ_MAP_HEADER)
562+
if write_bytes is not None:
563+
write_bytes(self._buffer[self._buf_checkpoint:self._buff_i])
560564
self._consume()
561565
return ret
562566

0 commit comments

Comments
 (0)