From db6f144ec0717c09b70cf07f9a9bd75f3c2851fe Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Mon, 26 Jun 2023 08:21:23 -0400
Subject: [PATCH] CLN: Make bz2 import optional

---
 pandas/compat/compressors.py | 30 +++++++++++++++++++-----------
 pandas/io/common.py          | 10 +++++++++-
 2 files changed, 28 insertions(+), 12 deletions(-)

diff --git a/pandas/compat/compressors.py b/pandas/compat/compressors.py
index a4f39c4e34bd42..1f31e34c092c96 100644
--- a/pandas/compat/compressors.py
+++ b/pandas/compat/compressors.py
@@ -4,11 +4,17 @@
 
 from __future__ import annotations
 
-import bz2
 from pickle import PickleBuffer
 
 from pandas.compat._constants import PY310
 
+try:
+    import bz2
+
+    has_bz2 = True
+except ImportError:
+    has_bz2 = False
+
 try:
     import lzma
 
@@ -41,17 +47,19 @@ def flatten_buffer(
         return memoryview(b).tobytes("A")
 
 
-class BZ2File(bz2.BZ2File):
-    if not PY310:
+if has_bz2:
 
-        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))
+    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:
diff --git a/pandas/io/common.py b/pandas/io/common.py
index 43780a08a4339c..d96f784b07c269 100644
--- a/pandas/io/common.py
+++ b/pandas/io/common.py
@@ -59,7 +59,6 @@
 )
 from pandas.compat import get_lzma_file
 from pandas.compat._optional import import_optional_dependency
-from pandas.compat.compressors import BZ2File as _BZ2File
 from pandas.util._decorators import doc
 from pandas.util._exceptions import find_stack_level
 
@@ -73,6 +72,11 @@
 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+\-+.]*://")
@@ -766,6 +770,10 @@ 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,
                 mode=ioargs.mode,