Move all compat compressors into a single file

Author: jakirkham

pandas/compat/__init__.py

@@ -12,9 +12,9 @@
 import os
 import platform
 import sys
-from typing import TYPE_CHECKING
 
 from pandas._typing import F
+import pandas.compat._compressors
 from pandas.compat.numpy import (
     is_numpy_dev,
     np_version_under1p21,
@@ -26,9 +26,6 @@
     pa_version_under9p0,
 )
 
-if TYPE_CHECKING:
-    import pandas.compat.lzma
-
 PY39 = sys.version_info >= (3, 9)
 PY310 = sys.version_info >= (3, 10)
 PY311 = sys.version_info >= (3, 11)
@@ -121,7 +118,7 @@ def is_ci_environment() -> bool:
     return os.environ.get("PANDAS_CI", "0") == "1"
 
 
-def get_lzma_file() -> type[pandas.compat.lzma.LZMAFile]:
+def get_lzma_file() -> type[pandas.compat._compressors.LZMAFile]:
     """
     Importing the `LZMAFile` class from the `lzma` module.
 
@@ -135,15 +132,13 @@ def get_lzma_file() -> type[pandas.compat.lzma.LZMAFile]:
     RuntimeError
         If the `lzma` module was not imported correctly, or didn't exist.
     """
-    try:
-        import pandas.compat.lzma
-    except ImportError:
+    if not pandas.compat._compressors.has_lzma:
         raise RuntimeError(
             "lzma module not available. "
             "A Python re-install with the proper dependencies, "
             "might be required to solve this issue."
         )
-    return pandas.compat.lzma.LZMAFile
+    return pandas.compat._compressors.LZMAFile
 
 
 __all__ = [

pandas/compat/_compressors.py

@@ -0,0 +1,78 @@
+"""
+Patched ``BZ2File`` and ``LZMAFile`` to handle pickle protocol 5.
+"""
+
+from __future__ import annotations
+
+import bz2
+import pickle
+import sys
+
+try:
+    import lzma
+
+    has_lzma = True
+except ImportError:
+    has_lzma = False
+
+
+PY310 = sys.version_info >= (3, 10)
+
+
+def flatten_buffer(
+    b: bytes | bytearray | memoryview | pickle.PickleBuffer,
+) -> bytes | 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, pickle.PickleBuffer):
+        b = pickle.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 not PY310:
+
+    class BZ2File(bz2.BZ2File):
+        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):
+            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))
+
+else:
+
+    class BZ2File(bz2.BZ2File):
+        pass
+
+    if has_lzma:
+
+        class LZMAFile(lzma.LZMAFile):
+            pass

pandas/compat/_utils.py

@@ -38,8 +38,8 @@
     get_lzma_file,
     is_platform_little_endian,
 )
+from pandas.compat._compressors import flatten_buffer
 from pandas.compat._optional import import_optional_dependency
-from pandas.compat._utils import flatten_buffer
 import pandas.util._test_decorators as td
 
 import pandas as pd