|
5 | 5 | of the parsers defined in parsers.py
|
6 | 6 | """
|
7 | 7 |
|
8 |
| -import bz2 |
9 |
| -import gzip |
10 | 8 | import os
|
11 | 9 | import zipfile
|
12 | 10 |
|
13 | 11 | import pytest
|
14 | 12 |
|
15 |
| -import pandas.compat as compat |
16 |
| - |
17 | 13 | import pandas as pd
|
18 | 14 | import pandas.util.testing as tm
|
19 | 15 |
|
20 | 16 |
|
21 |
| -def lzma_file(): |
22 |
| - """ |
23 |
| - Try to load the `LZMAFile` class from `backports.lzma`. |
24 |
| -
|
25 |
| - Returns |
26 |
| - ------- |
27 |
| - klass : type or None |
28 |
| - """ |
29 |
| - try: |
30 |
| - lzma = compat.import_lzma() |
31 |
| - except ImportError: |
32 |
| - lzma = None |
33 |
| - |
34 |
| - return getattr(lzma, "LZMAFile", None) |
35 |
| - |
36 |
| - |
37 |
| -def write_to_compressed(compress_type, path, data, dest="test"): |
38 |
| - """ |
39 |
| - Write data to a compressed file. |
40 |
| -
|
41 |
| - Parameters |
42 |
| - ---------- |
43 |
| - compress_type : type |
44 |
| - The compression type (or class) to use. |
45 |
| - path : str |
46 |
| - The file path to write the data. |
47 |
| - data : str |
48 |
| - The data to write. |
49 |
| - dest : str, default "test" |
50 |
| - The destination file (for ZIP only) |
51 |
| - """ |
52 |
| - # compression --> compression_method |
53 |
| - compression_mappings = { |
54 |
| - "zip": zipfile.ZipFile, |
55 |
| - "gzip": gzip.GzipFile, |
56 |
| - "bz2": bz2.BZ2File, |
57 |
| - "xz": lzma_file(), |
58 |
| - } |
59 |
| - |
60 |
| - compress_method = compression_mappings[compress_type] |
61 |
| - |
62 |
| - if compress_type == "zip": |
63 |
| - mode = "w" |
64 |
| - args = (dest, data) |
65 |
| - method = "writestr" |
66 |
| - else: |
67 |
| - mode = "wb" |
68 |
| - args = (data,) |
69 |
| - method = "write" |
70 |
| - |
71 |
| - with compress_method(path, mode=mode) as f: |
72 |
| - getattr(f, method)(*args) |
73 |
| - |
74 |
| - |
75 | 17 | @pytest.fixture(params=[True, False])
|
76 | 18 | def buffer(request):
|
77 | 19 | return request.param
|
@@ -154,7 +96,7 @@ def test_compression(parser_and_data, compression_only, buffer, filename):
|
154 | 96 | "buffer of compressed data.")
|
155 | 97 |
|
156 | 98 | with tm.ensure_clean(filename=filename) as path:
|
157 |
| - write_to_compressed(compress_type, path, data) |
| 99 | + tm.write_to_compressed(compress_type, path, data) |
158 | 100 | compression = "infer" if filename else compress_type
|
159 | 101 |
|
160 | 102 | if buffer:
|
|
0 commit comments