File size: 1,975 Bytes
6370773 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 |
"""
Patched ``BZ2File`` and ``LZMAFile`` to handle pickle protocol 5.
"""
from __future__ import annotations
from pickle import PickleBuffer
from pandas.compat._constants import PY310
try:
import bz2
has_bz2 = True
except ImportError:
has_bz2 = False
try:
import lzma
has_lzma = True
except ImportError:
has_lzma = False
def flatten_buffer(
b: bytes | bytearray | memoryview | PickleBuffer,
) -> bytes | bytearray | memoryview:
"""
Return some 1-D `uint8` typed buffer.
Coerces anything that does not match that description to one that does
without copying if possible (otherwise will copy).
"""
if isinstance(b, (bytes, bytearray)):
return b
if not isinstance(b, PickleBuffer):
b = PickleBuffer(b)
try:
# coerce to 1-D `uint8` C-contiguous `memoryview` zero-copy
return b.raw()
except BufferError:
# perform in-memory copy if buffer is not contiguous
return memoryview(b).tobytes("A")
if has_bz2:
class BZ2File(bz2.BZ2File):
if not PY310:
def write(self, b) -> int:
# Workaround issue where `bz2.BZ2File` expects `len`
# to return the number of bytes in `b` by converting
# `b` into something that meets that constraint with
# minimal copying.
#
# Note: This is fixed in Python 3.10.
return super().write(flatten_buffer(b))
if has_lzma:
class LZMAFile(lzma.LZMAFile):
if not PY310:
def write(self, b) -> int:
# Workaround issue where `lzma.LZMAFile` expects `len`
# to return the number of bytes in `b` by converting
# `b` into something that meets that constraint with
# minimal copying.
#
# Note: This is fixed in Python 3.10.
return super().write(flatten_buffer(b))
|