Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""
2Patched ``BZ2File`` and ``LZMAFile`` to handle pickle protocol 5.
3"""
5from __future__ import annotations
7from pickle import PickleBuffer
9from pandas.compat._constants import PY310
11try:
12 import bz2
14 has_bz2 = True
15except ImportError:
16 has_bz2 = False
18try:
19 import lzma
21 has_lzma = True
22except ImportError:
23 has_lzma = False
26def flatten_buffer(
27 b: bytes | bytearray | memoryview | PickleBuffer,
28) -> bytes | bytearray | memoryview:
29 """
30 Return some 1-D `uint8` typed buffer.
32 Coerces anything that does not match that description to one that does
33 without copying if possible (otherwise will copy).
34 """
36 if isinstance(b, (bytes, bytearray)):
37 return b
39 if not isinstance(b, PickleBuffer):
40 b = PickleBuffer(b)
42 try:
43 # coerce to 1-D `uint8` C-contiguous `memoryview` zero-copy
44 return b.raw()
45 except BufferError:
46 # perform in-memory copy if buffer is not contiguous
47 return memoryview(b).tobytes("A")
50if has_bz2:
52 class BZ2File(bz2.BZ2File):
53 if not PY310:
55 def write(self, b) -> int:
56 # Workaround issue where `bz2.BZ2File` expects `len`
57 # to return the number of bytes in `b` by converting
58 # `b` into something that meets that constraint with
59 # minimal copying.
60 #
61 # Note: This is fixed in Python 3.10.
62 return super().write(flatten_buffer(b))
65if has_lzma:
67 class LZMAFile(lzma.LZMAFile):
68 if not PY310:
70 def write(self, b) -> int:
71 # Workaround issue where `lzma.LZMAFile` expects `len`
72 # to return the number of bytes in `b` by converting
73 # `b` into something that meets that constraint with
74 # minimal copying.
75 #
76 # Note: This is fixed in Python 3.10.
77 return super().write(flatten_buffer(b))