CLN: Create get_bz2_file to match get_lzma_file

Author: MilesCranmer

pandas/_testing/_io.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import bz2
 from functools import wraps
 import gzip
 import io
@@ -13,7 +12,10 @@
 )
 import zipfile
 
-from pandas.compat import get_lzma_file
+from pandas.compat import (
+    get_bz2_file,
+    get_lzma_file,
+)
 from pandas.compat._optional import import_optional_dependency
 
 import pandas as pd
@@ -410,7 +412,7 @@ def write_to_compressed(compression, path, data, dest: str = "test"):
     elif compression == "gzip":
         compress_method = gzip.GzipFile
     elif compression == "bz2":
-        compress_method = bz2.BZ2File
+        compress_method = get_bz2_file()
     elif compression == "zstd":
         compress_method = import_optional_dependency("zstandard").open
     elif compression == "xz":

pandas/compat/__init__.py

@@ -154,6 +154,29 @@ def get_lzma_file() -> type[pandas.compat.compressors.LZMAFile]:
     return pandas.compat.compressors.LZMAFile
 
 
+def get_bz2_file() -> type[pandas.compat.compressors.BZ2File]:
+    """
+    Importing the `BZ2File` class from the `bz2` module.
+
+    Returns
+    -------
+    class
+        The `BZ2File` class from the `bz2` module.
+
+    Raises
+    ------
+    RuntimeError
+        If the `bz2` module was not imported correctly, or didn't exist.
+    """
+    if not pandas.compat.compressors.has_bz2:
+        raise RuntimeError(
+            "bz2 module not available. "
+            "A Python re-install with the proper dependencies, "
+            "might be required to solve this issue."
+        )
+    return pandas.compat.compressors.BZ2File
+
+
 __all__ = [
     "is_numpy_dev",
     "pa_version_under7p0",

pandas/io/common.py

@@ -57,7 +57,10 @@
     StorageOptions,
     WriteBuffer,
 )
-from pandas.compat import get_lzma_file
+from pandas.compat import (
+    get_bz2_file,
+    get_lzma_file,
+)
 from pandas.compat._optional import import_optional_dependency
 from pandas.util._decorators import doc
 from pandas.util._exceptions import find_stack_level
@@ -72,11 +75,6 @@
 from pandas.core.indexes.api import MultiIndex
 from pandas.core.shared_docs import _shared_docs
 
-try:
-    from pandas.compat.compressors import BZ2File as _BZ2File
-except ImportError:
-    _BZ2File = None
-
 _VALID_URLS = set(uses_relative + uses_netloc + uses_params)
 _VALID_URLS.discard("")
 _RFC_3986_PATTERN = re.compile(r"^[A-Za-z][A-Za-z0-9+\-+.]*://")
@@ -770,11 +768,7 @@ def get_handle(
         elif compression == "bz2":
             # Overload of "BZ2File" to handle pickle protocol 5
             # "Union[str, BaseBuffer]", "str", "Dict[str, Any]"
-            if _BZ2File is None:
-                raise ImportError(
-                    "bz2 compression requires the bz2 module to be installed"
-                )
-            handle = _BZ2File(  # type: ignore[call-overload]
+            handle = get_bz2_file()(  # type: ignore[call-overload]
                 handle,
                 mode=ioargs.mode,
                 **compression_args,