Skip to content

Commit 1e622c8

Browse files
twoertweinluckyvs1
authored andcommitted
CLN: simplify tm.decompress_file (pandas-dev#38758)
1 parent ea08083 commit 1e622c8

File tree

1 file changed

+3
-36
lines changed

1 file changed

+3
-36
lines changed

pandas/_testing/contexts.py

+3-36
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,9 @@
1-
import bz2
21
from contextlib import contextmanager
3-
import gzip
42
import os
53
from shutil import rmtree
64
import tempfile
7-
import zipfile
85

9-
from pandas.compat import get_lzma_file, import_lzma
10-
11-
lzma = import_lzma()
6+
from pandas.io.common import get_handle
127

138

149
@contextmanager
@@ -28,36 +23,8 @@ def decompress_file(path, compression):
2823
-------
2924
file object
3025
"""
31-
if compression is None:
32-
f = open(path, "rb")
33-
elif compression == "gzip":
34-
# pandas\_testing.py:243: error: Incompatible types in assignment
35-
# (expression has type "IO[Any]", variable has type "BinaryIO")
36-
f = gzip.open(path, "rb") # type: ignore[assignment]
37-
elif compression == "bz2":
38-
# pandas\_testing.py:245: error: Incompatible types in assignment
39-
# (expression has type "BZ2File", variable has type "BinaryIO")
40-
f = bz2.BZ2File(path, "rb") # type: ignore[assignment]
41-
elif compression == "xz":
42-
f = get_lzma_file(lzma)(path, "rb")
43-
elif compression == "zip":
44-
zip_file = zipfile.ZipFile(path)
45-
zip_names = zip_file.namelist()
46-
if len(zip_names) == 1:
47-
# pandas\_testing.py:252: error: Incompatible types in assignment
48-
# (expression has type "IO[bytes]", variable has type "BinaryIO")
49-
f = zip_file.open(zip_names.pop()) # type: ignore[assignment]
50-
else:
51-
raise ValueError(f"ZIP file {path} error. Only one file per ZIP.")
52-
else:
53-
raise ValueError(f"Unrecognized compression type: {compression}")
54-
55-
try:
56-
yield f
57-
finally:
58-
f.close()
59-
if compression == "zip":
60-
zip_file.close()
26+
with get_handle(path, "rb", compression=compression, is_text=False) as handle:
27+
yield handle.handle
6128

6229

6330
@contextmanager

0 commit comments

Comments
 (0)