codekingpro/portable-devtools
115k
1"""Fallback pure Python implementation of msgpack"""
2
3import struct
4import sys
5from datetime import datetime as _DateTime
6
7if hasattr(sys, "pypy_version_info"):
8 from __pypy__ import newlist_hint
9 from __pypy__.builders import BytesBuilder
10
11 _USING_STRINGBUILDER = True
12
13 class BytesIO:
14 def __init__(self, s=b""):
15 if s:
16 self.builder = BytesBuilder(len(s))
17 self.builder.append(s)
18 else:
19 self.builder = BytesBuilder()
20
21 def write(self, s):
22 if isinstance(s, memoryview):
23 s = s.tobytes()
24 elif isinstance(s, bytearray):
25 s = bytes(s)
26 self.builder.append(s)
27
28 def getvalue(self):
29 return self.builder.build()
30
31else:
32 from io import BytesIO
33
34 _USING_STRINGBUILDER = False
35
36 def newlist_hint(size):
37 return []
38
39
40from .exceptions import BufferFull, ExtraData, FormatError, OutOfData, StackError
41from .ext import ExtType, Timestamp
42
43EX_SKIP = 0
44EX_CONSTRUCT = 1
45EX_READ_ARRAY_HEADER = 2
46EX_READ_MAP_HEADER = 3
47
48TYPE_IMMEDIATE = 0
49TYPE_ARRAY = 1
50TYPE_MAP = 2
51TYPE_RAW = 3
52TYPE_BIN = 4
53TYPE_EXT = 5
54
55DEFAULT_RECURSE_LIMIT = 511
56
57
58def _check_type_strict(obj, t, type=type, tuple=tuple):
59 if type(t) is tuple:
60 return type(obj) in t
61 else:
62 return type(obj) is t
63
64
65def _get_data_from_buffer(obj):
66 view = memoryview(obj)
67 if view.itemsize != 1:
68 raise ValueError("cannot unpack from multi-byte object")
69 return view
70
71
72def unpackb(packed, **kwargs):
73 """
74 Unpack an object from `packed`.
75
76 Raises ``ExtraData`` when *packed* contains extra bytes.
77 Raises ``ValueError`` when *packed* is incomplete.
78 Raises ``FormatError`` when *packed* is not valid msgpack.
79 Raises ``StackError`` when *packed* contains too nested.
80 Other exceptions can be raised during unpacking.
81
82 See :class:`Unpacker` for options.
83 """
84 unpacker = Unpacker(None, max_buffer_size=len(packed), **kwargs)
85 unpacker.feed(packed)
86 try:
87 ret = unpacker._unpack()
88 except OutOfData:
89 raise ValueError("Unpack failed: incomplete input")
90 except RecursionError:
91 raise StackError
92 if unpacker._got_extradata():
93 raise ExtraData(ret, unpacker._get_extradata())
94 return ret
95
96
97_NO_FORMAT_USED = ""
98_MSGPACK_HEADERS = {
99 0xC4: (1, _NO_FORMAT_USED, TYPE_BIN),
100 0xC5: (2, ">H", TYPE_BIN),
101 0xC6: (4, ">I", TYPE_BIN),
102 0xC7: (2, "Bb", TYPE_EXT),
103 0xC8: (3, ">Hb", TYPE_EXT),
104 0xC9: (5, ">Ib", TYPE_EXT),
105 0xCA: (4, ">f"),
106 0xCB: (8, ">d"),
107 0xCC: (1, _NO_FORMAT_USED),
108 0xCD: (2, ">H"),
109 0xCE: (4, ">I"),
110 0xCF: (8, ">Q"),
111 0xD0: (1, "b"),
112 0xD1: (2, ">h"),
113 0xD2: (4, ">i"),
114 0xD3: (8, ">q"),
115 0xD4: (1, "b1s", TYPE_EXT),
116 0xD5: (2, "b2s", TYPE_EXT),
117 0xD6: (4, "b4s", TYPE_EXT),
118 0xD7: (8, "b8s", TYPE_EXT),
119 0xD8: (16, "b16s", TYPE_EXT),
120 0xD9: (1, _NO_FORMAT_USED, TYPE_RAW),
121 0xDA: (2, ">H", TYPE_RAW),
122 0xDB: (4, ">I", TYPE_RAW),
123 0xDC: (2, ">H", TYPE_ARRAY),
124 0xDD: (4, ">I", TYPE_ARRAY),
125 0xDE: (2, ">H", TYPE_MAP),
126 0xDF: (4, ">I", TYPE_MAP),
127}
128
129
130class Unpacker:
131 """Streaming unpacker.
132
133 Arguments:
134
135 :param file_like:
136 File-like object having `.read(n)` method.
137 If specified, unpacker reads serialized data from it and `.feed()` is not usable.
138
139 :param int read_size:
140 Used as `file_like.read(read_size)`. (default: `min(16*1024, max_buffer_size)`)
141
142 :param bool use_list:
143 If true, unpack msgpack array to Python list.
144 Otherwise, unpack to Python tuple. (default: True)
145
146 :param bool raw:
147 If true, unpack msgpack raw to Python bytes.
148 Otherwise, unpack to Python str by decoding with UTF-8 encoding (default).
149
150 :param int timestamp:
151 Control how timestamp type is unpacked:
152
153 0 - Timestamp
154 1 - float (Seconds from the EPOCH)
155 2 - int (Nanoseconds from the EPOCH)
156 3 - datetime.datetime (UTC).
157
158 :param bool strict_map_key:
159 If true (default), only str or bytes are accepted for map (dict) keys.
160
161 :param object_hook:
162 When specified, it should be callable.
163 Unpacker calls it with a dict argument after unpacking msgpack map.
164 (See also simplejson)
165
166 :param object_pairs_hook:
167 When specified, it should be callable.
168 Unpacker calls it with a list of key-value pairs after unpacking msgpack map.
169 (See also simplejson)
170
171 :param str unicode_errors:
172 The error handler for decoding unicode. (default: 'strict')
173 This option should be used only when you have msgpack data which
174 contains invalid UTF-8 string.
175
176 :param int max_buffer_size:
177 Limits size of data waiting unpacked. 0 means 2**32-1.
178 The default value is 100*1024*1024 (100MiB).
179 Raises `BufferFull` exception when it is insufficient.
180 You should set this parameter when unpacking data from untrusted source.
181
182 :param int max_str_len:
183 Deprecated, use *max_buffer_size* instead.
184 Limits max length of str. (default: max_buffer_size)
185
186 :param int max_bin_len:
187 Deprecated, use *max_buffer_size* instead.
188 Limits max length of bin. (default: max_buffer_size)
189
190 :param int max_array_len:
191 Limits max length of array.
192 (default: max_buffer_size)
193
194 :param int max_map_len:
195 Limits max length of map.
196 (default: max_buffer_size//2)
197
198 :param int max_ext_len:
199 Deprecated, use *max_buffer_size* instead.
200 Limits max size of ext type. (default: max_buffer_size)
201
202 Example of streaming deserialize from file-like object::
203
204 unpacker = Unpacker(file_like)
205 for o in unpacker:
206 process(o)
207
208 Example of streaming deserialize from socket::
209
210 unpacker = Unpacker()
211 while True:
212 buf = sock.recv(1024**2)
213 if not buf:
214 break
215 unpacker.feed(buf)
216 for o in unpacker:
217 process(o)
218
219 Raises ``ExtraData`` when *packed* contains extra bytes.
220 Raises ``OutOfData`` when *packed* is incomplete.
221 Raises ``FormatError`` when *packed* is not valid msgpack.
222 Raises ``StackError`` when *packed* contains too nested.
223 Other exceptions can be raised during unpacking.
224 """
225
226 def __init__(
227 self,
228 file_like=None,
229 *,
230 read_size=0,
231 use_list=True,
232 raw=False,
233 timestamp=0,
234 strict_map_key=True,
235 object_hook=None,
236 object_pairs_hook=None,
237 list_hook=None,
238 unicode_errors=None,
239 max_buffer_size=100 * 1024 * 1024,
240 ext_hook=ExtType,
241 max_str_len=-1,
242 max_bin_len=-1,
243 max_array_len=-1,
244 max_map_len=-1,
245 max_ext_len=-1,
246 ):
247 if unicode_errors is None:
248 unicode_errors = "strict"
249
250 if file_like is None:
251 self._feeding = True
252 else:
253 if not callable(file_like.read):
254 raise TypeError("`file_like.read` must be callable")
255 self.file_like = file_like
256 self._feeding = False
257
258 #: array of bytes fed.
259 self._buffer = bytearray()
260 #: Which position we currently reads
261 self._buff_i = 0
262
263 # When Unpacker is used as an iterable, between the calls to next(),
264 # the buffer is not "consumed" completely, for efficiency sake.
265 # Instead, it is done sloppily. To make sure we raise BufferFull at
266 # the correct moments, we have to keep track of how sloppy we were.
267 # Furthermore, when the buffer is incomplete (that is: in the case
268 # we raise an OutOfData) we need to rollback the buffer to the correct
269 # state, which _buf_checkpoint records.
270 self._buf_checkpoint = 0
271
272 if not max_buffer_size:
273 max_buffer_size = 2**31 - 1
274 if max_str_len == -1:
275 max_str_len = max_buffer_size
276 if max_bin_len == -1:
277 max_bin_len = max_buffer_size
278 if max_array_len == -1:
279 max_array_len = max_buffer_size
280 if max_map_len == -1:
281 max_map_len = max_buffer_size // 2
282 if max_ext_len == -1:
283 max_ext_len = max_buffer_size
284
285 self._max_buffer_size = max_buffer_size
286 if read_size > self._max_buffer_size:
287 raise ValueError("read_size must be smaller than max_buffer_size")
288 self._read_size = read_size or min(self._max_buffer_size, 16 * 1024)
289 self._raw = bool(raw)
290 self._strict_map_key = bool(strict_map_key)
291 self._unicode_errors = unicode_errors
292 self._use_list = use_list
293 if not (0 <= timestamp <= 3):
294 raise ValueError("timestamp must be 0..3")
295 self._timestamp = timestamp
296 self._list_hook = list_hook
297 self._object_hook = object_hook
298 self._object_pairs_hook = object_pairs_hook
299 self._ext_hook = ext_hook
300 self._max_str_len = max_str_len
301 self._max_bin_len = max_bin_len
302 self._max_array_len = max_array_len
303 self._max_map_len = max_map_len
304 self._max_ext_len = max_ext_len
305 self._stream_offset = 0
306
307 if list_hook is not None and not callable(list_hook):
308 raise TypeError("`list_hook` is not callable")
309 if object_hook is not None and not callable(object_hook):
310 raise TypeError("`object_hook` is not callable")
311 if object_pairs_hook is not None and not callable(object_pairs_hook):
312 raise TypeError("`object_pairs_hook` is not callable")
313 if object_hook is not None and object_pairs_hook is not None:
314 raise TypeError("object_pairs_hook and object_hook are mutually exclusive")
315 if not callable(ext_hook):
316 raise TypeError("`ext_hook` is not callable")
317
318 def feed(self, next_bytes):
319 assert self._feeding
320 view = _get_data_from_buffer(next_bytes)
321 if len(self._buffer) - self._buff_i + len(view) > self._max_buffer_size:
322 raise BufferFull
323
324 # Strip buffer before checkpoint before reading file.
325 if self._buf_checkpoint > 0:
326 del self._buffer[: self._buf_checkpoint]
327 self._buff_i -= self._buf_checkpoint
328 self._buf_checkpoint = 0
329
330 # Use extend here: INPLACE_ADD += doesn't reliably typecast memoryview in jython
331 self._buffer.extend(view)
332 view.release()
333
334 def _consume(self):
335 """Gets rid of the used parts of the buffer."""
336 self._stream_offset += self._buff_i - self._buf_checkpoint
337 self._buf_checkpoint = self._buff_i
338
339 def _got_extradata(self):
340 return self._buff_i < len(self._buffer)
341
342 def _get_extradata(self):
343 return self._buffer[self._buff_i :]
344
345 def read_bytes(self, n):
346 ret = self._read(n, raise_outofdata=False)
347 self._consume()
348 return ret
349
350 def _read(self, n, raise_outofdata=True):
351 # (int) -> bytearray
352 self._reserve(n, raise_outofdata=raise_outofdata)
353 i = self._buff_i
354 ret = self._buffer[i : i + n]
355 self._buff_i = i + len(ret)
356 return ret
357
358 def _reserve(self, n, raise_outofdata=True):
359 remain_bytes = len(self._buffer) - self._buff_i - n
360
361 # Fast path: buffer has n bytes already
362 if remain_bytes >= 0:
363 return
364
365 if self._feeding:
366 self._buff_i = self._buf_checkpoint
367 raise OutOfData
368
369 # Strip buffer before checkpoint before reading file.
370 if self._buf_checkpoint > 0:
371 del self._buffer[: self._buf_checkpoint]
372 self._buff_i -= self._buf_checkpoint
373 self._buf_checkpoint = 0
374
375 # Read from file
376 remain_bytes = -remain_bytes
377 if remain_bytes + len(self._buffer) > self._max_buffer_size:
378 raise BufferFull
379 while remain_bytes > 0:
380 to_read_bytes = max(self._read_size, remain_bytes)
381 read_data = self.file_like.read(to_read_bytes)
382 if not read_data:
383 break
384 assert isinstance(read_data, bytes)
385 self._buffer += read_data
386 remain_bytes -= len(read_data)
387
388 if len(self._buffer) < n + self._buff_i and raise_outofdata:
389 self._buff_i = 0 # rollback
390 raise OutOfData
391
392 def _read_header(self):
393 typ = TYPE_IMMEDIATE
394 n = 0
395 obj = None
396 self._reserve(1)
397 b = self._buffer[self._buff_i]
398 self._buff_i += 1
399 if b & 0b10000000 == 0:
400 obj = b
401 elif b & 0b11100000 == 0b11100000:
402 obj = -1 - (b ^ 0xFF)
403 elif b & 0b11100000 == 0b10100000:
404 n = b & 0b00011111
405 typ = TYPE_RAW
406 if n > self._max_str_len:
407 raise ValueError(f"{n} exceeds max_str_len({self._max_str_len})")
408 obj = self._read(n)
409 elif b & 0b11110000 == 0b10010000:
410 n = b & 0b00001111
411 typ = TYPE_ARRAY
412 if n > self._max_array_len:
413 raise ValueError(f"{n} exceeds max_array_len({self._max_array_len})")
414 elif b & 0b11110000 == 0b10000000:
415 n = b & 0b00001111
416 typ = TYPE_MAP
417 if n > self._max_map_len:
418 raise ValueError(f"{n} exceeds max_map_len({self._max_map_len})")
419 elif b == 0xC0:
420 obj = None
421 elif b == 0xC2:
422 obj = False
423 elif b == 0xC3:
424 obj = True
425 elif 0xC4 <= b <= 0xC6:
426 size, fmt, typ = _MSGPACK_HEADERS[b]
427 self._reserve(size)
428 if len(fmt) > 0:
429 n = struct.unpack_from(fmt, self._buffer, self._buff_i)[0]
430 else:
431 n = self._buffer[self._buff_i]
432 self._buff_i += size
433 if n > self._max_bin_len:
434 raise ValueError(f"{n} exceeds max_bin_len({self._max_bin_len})")
435 obj = self._read(n)
436 elif 0xC7 <= b <= 0xC9:
437 size, fmt, typ = _MSGPACK_HEADERS[b]
438 self._reserve(size)
439 L, n = struct.unpack_from(fmt, self._buffer, self._buff_i)
440 self._buff_i += size
441 if L > self._max_ext_len:
442 raise ValueError(f"{L} exceeds max_ext_len({self._max_ext_len})")
443 obj = self._read(L)
444 elif 0xCA <= b <= 0xD3:
445 size, fmt = _MSGPACK_HEADERS[b]
446 self._reserve(size)
447 if len(fmt) > 0:
448 obj = struct.unpack_from(fmt, self._buffer, self._buff_i)[0]
449 else:
450 obj = self._buffer[self._buff_i]
451 self._buff_i += size
452 elif 0xD4 <= b <= 0xD8:
453 size, fmt, typ = _MSGPACK_HEADERS[b]
454 if self._max_ext_len < size:
455 raise ValueError(f"{size} exceeds max_ext_len({self._max_ext_len})")
456 self._reserve(size + 1)
457 n, obj = struct.unpack_from(fmt, self._buffer, self._buff_i)
458 self._buff_i += size + 1
459 elif 0xD9 <= b <= 0xDB:
460 size, fmt, typ = _MSGPACK_HEADERS[b]
461 self._reserve(size)
462 if len(fmt) > 0:
463 (n,) = struct.unpack_from(fmt, self._buffer, self._buff_i)
464 else:
465 n = self._buffer[self._buff_i]
466 self._buff_i += size
467 if n > self._max_str_len:
468 raise ValueError(f"{n} exceeds max_str_len({self._max_str_len})")
469 obj = self._read(n)
470 elif 0xDC <= b <= 0xDD:
471 size, fmt, typ = _MSGPACK_HEADERS[b]
472 self._reserve(size)
473 (n,) = struct.unpack_from(fmt, self._buffer, self._buff_i)
474 self._buff_i += size
475 if n > self._max_array_len:
476 raise ValueError(f"{n} exceeds max_array_len({self._max_array_len})")
477 elif 0xDE <= b <= 0xDF:
478 size, fmt, typ = _MSGPACK_HEADERS[b]
479 self._reserve(size)
480 (n,) = struct.unpack_from(fmt, self._buffer, self._buff_i)
481 self._buff_i += size
482 if n > self._max_map_len:
483 raise ValueError(f"{n} exceeds max_map_len({self._max_map_len})")
484 else:
485 raise FormatError("Unknown header: 0x%x" % b)
486 return typ, n, obj
487
488 def _unpack(self, execute=EX_CONSTRUCT):
489 typ, n, obj = self._read_header()
490
491 if execute == EX_READ_ARRAY_HEADER:
492 if typ != TYPE_ARRAY:
493 raise ValueError("Expected array")
494 return n
495 if execute == EX_READ_MAP_HEADER:
496 if typ != TYPE_MAP:
497 raise ValueError("Expected map")
498 return n
499 # TODO should we eliminate the recursion?
500 if typ == TYPE_ARRAY:
501 if execute == EX_SKIP:
502 for i in range(n):
503 # TODO check whether we need to call `list_hook`
504 self._unpack(EX_SKIP)
505 return
506 ret = newlist_hint(n)
507 for i in range(n):
508 ret.append(self._unpack(EX_CONSTRUCT))
509 if self._list_hook is not None:
510 ret = self._list_hook(ret)
511 # TODO is the interaction between `list_hook` and `use_list` ok?
512 return ret if self._use_list else tuple(ret)
513 if typ == TYPE_MAP:
514 if execute == EX_SKIP:
515 for i in range(n):
516 # TODO check whether we need to call hooks
517 self._unpack(EX_SKIP)
518 self._unpack(EX_SKIP)
519 return
520 if self._object_pairs_hook is not None:
521 ret = self._object_pairs_hook(
522 (self._unpack(EX_CONSTRUCT), self._unpack(EX_CONSTRUCT)) for _ in range(n)
523 )
524 else:
525 ret = {}
526 for _ in range(n):
527 key = self._unpack(EX_CONSTRUCT)
528 if self._strict_map_key and type(key) not in (str, bytes):
529 raise ValueError("%s is not allowed for map key" % str(type(key)))
530 if isinstance(key, str):
531 key = sys.intern(key)
532 ret[key] = self._unpack(EX_CONSTRUCT)
533 if self._object_hook is not None:
534 ret = self._object_hook(ret)
535 return ret
536 if execute == EX_SKIP:
537 return
538 if typ == TYPE_RAW:
539 if self._raw:
540 obj = bytes(obj)
541 else:
542 obj = obj.decode("utf_8", self._unicode_errors)
543 return obj
544 if typ == TYPE_BIN:
545 return bytes(obj)
546 if typ == TYPE_EXT:
547 if n == -1: # timestamp
548 ts = Timestamp.from_bytes(bytes(obj))
549 if self._timestamp == 1:
550 return ts.to_unix()
551 elif self._timestamp == 2:
552 return ts.to_unix_nano()
553 elif self._timestamp == 3:
554 return ts.to_datetime()
555 else:
556 return ts
557 else:
558 return self._ext_hook(n, bytes(obj))
559 assert typ == TYPE_IMMEDIATE
560 return obj
561
562 def __iter__(self):
563 return self
564
565 def __next__(self):
566 try:
567 ret = self._unpack(EX_CONSTRUCT)
568 self._consume()
569 return ret
570 except OutOfData:
571 self._consume()
572 raise StopIteration
573 except RecursionError:
574 raise StackError
575
576 next = __next__
577
578 def skip(self):
579 self._unpack(EX_SKIP)
580 self._consume()
581
582 def unpack(self):
583 try:
584 ret = self._unpack(EX_CONSTRUCT)
585 except RecursionError:
586 raise StackError
587 self._consume()
588 return ret
589
590 def read_array_header(self):
591 ret = self._unpack(EX_READ_ARRAY_HEADER)
592 self._consume()
593 return ret
594
595 def read_map_header(self):
596 ret = self._unpack(EX_READ_MAP_HEADER)
597 self._consume()
598 return ret
599
600 def tell(self):
601 return self._stream_offset
602
603
604class Packer:
605 """
606 MessagePack Packer
607
608 Usage::
609
610 packer = Packer()
611 astream.write(packer.pack(a))
612 astream.write(packer.pack(b))
613
614 Packer's constructor has some keyword arguments:
615
616 :param default:
617 When specified, it should be callable.
618 Convert user type to builtin type that Packer supports.
619 See also simplejson's document.
620
621 :param bool use_single_float:
622 Use single precision float type for float. (default: False)
623
624 :param bool autoreset:
625 Reset buffer after each pack and return its content as `bytes`. (default: True).
626 If set this to false, use `bytes()` to get content and `.reset()` to clear buffer.
627
628 :param bool use_bin_type:
629 Use bin type introduced in msgpack spec 2.0 for bytes.
630 It also enables str8 type for unicode. (default: True)
631
632 :param bool strict_types:
633 If set to true, types will be checked to be exact. Derived classes
634 from serializable types will not be serialized and will be
635 treated as unsupported type and forwarded to default.
636 Additionally tuples will not be serialized as lists.
637 This is useful when trying to implement accurate serialization
638 for python types.
639
640 :param bool datetime:
641 If set to true, datetime with tzinfo is packed into Timestamp type.
642 Note that the tzinfo is stripped in the timestamp.
643 You can get UTC datetime with `timestamp=3` option of the Unpacker.
644
645 :param str unicode_errors:
646 The error handler for encoding unicode. (default: 'strict')
647 DO NOT USE THIS!! This option is kept for very specific usage.
648
649 :param int buf_size:
650 Internal buffer size. This option is used only for C implementation.
651 """
652
653 def __init__(
654 self,
655 *,
656 default=None,
657 use_single_float=False,
658 autoreset=True,
659 use_bin_type=True,
660 strict_types=False,
661 datetime=False,
662 unicode_errors=None,
663 buf_size=None,
664 ):
665 self._strict_types = strict_types
666 self._use_float = use_single_float
667 self._autoreset = autoreset
668 self._use_bin_type = use_bin_type
669 self._buffer = BytesIO()
670 self._datetime = bool(datetime)
671 self._unicode_errors = unicode_errors or "strict"
672 if default is not None and not callable(default):
673 raise TypeError("default must be callable")
674 self._default = default
675
676 def _pack(
677 self,
678 obj,
679 nest_limit=DEFAULT_RECURSE_LIMIT,
680 check=isinstance,
681 check_type_strict=_check_type_strict,
682 ):
683 default_used = False
684 if self._strict_types:
685 check = check_type_strict
686 list_types = list
687 else:
688 list_types = (list, tuple)
689 while True:
690 if nest_limit < 0:
691 raise ValueError("recursion limit exceeded")
692 if obj is None:
693 return self._buffer.write(b"\xc0")
694 if check(obj, bool):
695 if obj:
696 return self._buffer.write(b"\xc3")
697 return self._buffer.write(b"\xc2")
698 if check(obj, int):
699 if 0 <= obj < 0x80:
700 return self._buffer.write(struct.pack("B", obj))
701 if -0x20 <= obj < 0:
702 return self._buffer.write(struct.pack("b", obj))
703 if 0x80 <= obj <= 0xFF:
704 return self._buffer.write(struct.pack("BB", 0xCC, obj))
705 if -0x80 <= obj < 0:
706 return self._buffer.write(struct.pack(">Bb", 0xD0, obj))
707 if 0xFF < obj <= 0xFFFF:
708 return self._buffer.write(struct.pack(">BH", 0xCD, obj))
709 if -0x8000 <= obj < -0x80:
710 return self._buffer.write(struct.pack(">Bh", 0xD1, obj))
711 if 0xFFFF < obj <= 0xFFFFFFFF:
712 return self._buffer.write(struct.pack(">BI", 0xCE, obj))
713 if -0x80000000 <= obj < -0x8000:
714 return self._buffer.write(struct.pack(">Bi", 0xD2, obj))
715 if 0xFFFFFFFF < obj <= 0xFFFFFFFFFFFFFFFF:
716 return self._buffer.write(struct.pack(">BQ", 0xCF, obj))
717 if -0x8000000000000000 <= obj < -0x80000000:
718 return self._buffer.write(struct.pack(">Bq", 0xD3, obj))
719 if not default_used and self._default is not None:
720 obj = self._default(obj)
721 default_used = True
722 continue
723 raise OverflowError("Integer value out of range")
724 if check(obj, (bytes, bytearray)):
725 n = len(obj)
726 if n >= 2**32:
727 raise ValueError("%s is too large" % type(obj).__name__)
728 self._pack_bin_header(n)
729 return self._buffer.write(obj)
730 if check(obj, str):
731 obj = obj.encode("utf-8", self._unicode_errors)
732 n = len(obj)
733 if n >= 2**32:
734 raise ValueError("String is too large")
735 self._pack_raw_header(n)
736 return self._buffer.write(obj)
737 if check(obj, memoryview):
738 n = obj.nbytes
739 if n >= 2**32:
740 raise ValueError("Memoryview is too large")
741 self._pack_bin_header(n)
742 return self._buffer.write(obj)
743 if check(obj, float):
744 if self._use_float:
745 return self._buffer.write(struct.pack(">Bf", 0xCA, obj))
746 return self._buffer.write(struct.pack(">Bd", 0xCB, obj))
747 if check(obj, (ExtType, Timestamp)):
748 if check(obj, Timestamp):
749 code = -1
750 data = obj.to_bytes()
751 else:
752 code = obj.code
753 data = obj.data
754 assert isinstance(code, int)
755 assert isinstance(data, bytes)
756 L = len(data)
757 if L == 1:
758 self._buffer.write(b"\xd4")
759 elif L == 2:
760 self._buffer.write(b"\xd5")
761 elif L == 4:
762 self._buffer.write(b"\xd6")
763 elif L == 8:
764 self._buffer.write(b"\xd7")
765 elif L == 16:
766 self._buffer.write(b"\xd8")
767 elif L <= 0xFF:
768 self._buffer.write(struct.pack(">BB", 0xC7, L))
769 elif L <= 0xFFFF:
770 self._buffer.write(struct.pack(">BH", 0xC8, L))
771 else:
772 self._buffer.write(struct.pack(">BI", 0xC9, L))
773 self._buffer.write(struct.pack("b", code))
774 self._buffer.write(data)
775 return
776 if check(obj, list_types):
777 n = len(obj)
778 self._pack_array_header(n)
779 for i in range(n):
780 self._pack(obj[i], nest_limit - 1)
781 return
782 if check(obj, dict):
783 return self._pack_map_pairs(len(obj), obj.items(), nest_limit - 1)
784
785 if self._datetime and check(obj, _DateTime) and obj.tzinfo is not None:
786 obj = Timestamp.from_datetime(obj)
787 default_used = 1
788 continue
789
790 if not default_used and self._default is not None:
791 obj = self._default(obj)
792 default_used = 1
793 continue
794
795 if self._datetime and check(obj, _DateTime):
796 raise ValueError(f"Cannot serialize {obj!r} where tzinfo=None")
797
798 raise TypeError(f"Cannot serialize {obj!r}")
799
800 def pack(self, obj):
801 try:
802 self._pack(obj)
803 except:
804 self._buffer = BytesIO() # force reset
805 raise
806 if self._autoreset:
807 ret = self._buffer.getvalue()
808 self._buffer = BytesIO()
809 return ret
810
811 def pack_map_pairs(self, pairs):
812 self._pack_map_pairs(len(pairs), pairs)
813 if self._autoreset:
814 ret = self._buffer.getvalue()
815 self._buffer = BytesIO()
816 return ret
817
818 def pack_array_header(self, n):
819 if n >= 2**32:
820 raise ValueError
821 self._pack_array_header(n)
822 if self._autoreset:
823 ret = self._buffer.getvalue()
824 self._buffer = BytesIO()
825 return ret
826
827 def pack_map_header(self, n):
828 if n >= 2**32:
829 raise ValueError
830 self._pack_map_header(n)
831 if self._autoreset:
832 ret = self._buffer.getvalue()
833 self._buffer = BytesIO()
834 return ret
835
836 def pack_ext_type(self, typecode, data):
837 if not isinstance(typecode, int):
838 raise TypeError("typecode must have int type.")
839 if not 0 <= typecode <= 127:
840 raise ValueError("typecode should be 0-127")
841 if not isinstance(data, bytes):
842 raise TypeError("data must have bytes type")
843 L = len(data)
844 if L > 0xFFFFFFFF:
845 raise ValueError("Too large data")
846 if L == 1:
847 self._buffer.write(b"\xd4")
848 elif L == 2:
849 self._buffer.write(b"\xd5")
850 elif L == 4:
851 self._buffer.write(b"\xd6")
852 elif L == 8:
853 self._buffer.write(b"\xd7")
854 elif L == 16:
855 self._buffer.write(b"\xd8")
856 elif L <= 0xFF:
857 self._buffer.write(b"\xc7" + struct.pack("B", L))
858 elif L <= 0xFFFF:
859 self._buffer.write(b"\xc8" + struct.pack(">H", L))
860 else:
861 self._buffer.write(b"\xc9" + struct.pack(">I", L))
862 self._buffer.write(struct.pack("B", typecode))
863 self._buffer.write(data)
864
865 def _pack_array_header(self, n):
866 if n <= 0x0F:
867 return self._buffer.write(struct.pack("B", 0x90 + n))
868 if n <= 0xFFFF:
869 return self._buffer.write(struct.pack(">BH", 0xDC, n))
870 if n <= 0xFFFFFFFF:
871 return self._buffer.write(struct.pack(">BI", 0xDD, n))
872 raise ValueError("Array is too large")
873
874 def _pack_map_header(self, n):
875 if n <= 0x0F:
876 return self._buffer.write(struct.pack("B", 0x80 + n))
877 if n <= 0xFFFF:
878 return self._buffer.write(struct.pack(">BH", 0xDE, n))
879 if n <= 0xFFFFFFFF:
880 return self._buffer.write(struct.pack(">BI", 0xDF, n))
881 raise ValueError("Dict is too large")
882
883 def _pack_map_pairs(self, n, pairs, nest_limit=DEFAULT_RECURSE_LIMIT):
884 self._pack_map_header(n)
885 for k, v in pairs:
886 self._pack(k, nest_limit - 1)
887 self._pack(v, nest_limit - 1)
888
889 def _pack_raw_header(self, n):
890 if n <= 0x1F:
891 self._buffer.write(struct.pack("B", 0xA0 + n))
892 elif self._use_bin_type and n <= 0xFF:
893 self._buffer.write(struct.pack(">BB", 0xD9, n))
894 elif n <= 0xFFFF:
895 self._buffer.write(struct.pack(">BH", 0xDA, n))
896 elif n <= 0xFFFFFFFF:
897 self._buffer.write(struct.pack(">BI", 0xDB, n))
898 else:
899 raise ValueError("Raw is too large")
900
901 def _pack_bin_header(self, n):
902 if not self._use_bin_type:
903 return self._pack_raw_header(n)
904 elif n <= 0xFF:
905 return self._buffer.write(struct.pack(">BB", 0xC4, n))
906 elif n <= 0xFFFF:
907 return self._buffer.write(struct.pack(">BH", 0xC5, n))
908 elif n <= 0xFFFFFFFF:
909 return self._buffer.write(struct.pack(">BI", 0xC6, n))
910 else:
911 raise ValueError("Bin is too large")
912
913 def bytes(self):
914 """Return internal buffer contents as bytes object"""
915 return self._buffer.getvalue()
916
917 def reset(self):
918 """Reset internal buffer.
919
920 This method is useful only when autoreset=False.
921 """
922 self._buffer = BytesIO()
923
924 def getbuffer(self):
925 """Return view of internal buffer."""
926 if _USING_STRINGBUILDER:
927 return memoryview(self.bytes())
928 else:
929 return self._buffer.getbuffer()
930