Skip to content

Commit

Permalink
CLN: Make bz2 import optional
Browse files Browse the repository at this point in the history
  • Loading branch information
MilesCranmer committed Jun 26, 2023
1 parent 49851bb commit db6f144
Show file tree
Hide file tree
Showing 2 changed files with 28 additions and 12 deletions.
30 changes: 19 additions & 11 deletions pandas/compat/compressors.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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:
Expand Down
10 changes: 9 additions & 1 deletion pandas/io/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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+\-+.]*://")
Expand Down Expand Up @@ -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,
Expand Down

0 comments on commit db6f144

Please sign in to comment.