-
Notifications
You must be signed in to change notification settings - Fork 238
Expand file tree
/
Copy pathcompress.py
More file actions
97 lines (72 loc) · 2.38 KB
/
compress.py
File metadata and controls
97 lines (72 loc) · 2.38 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
from typing import IO, TYPE_CHECKING, Callable, Optional
from docarray.utils._internal.misc import import_library
def _compress_bytes(data: bytes, algorithm: Optional[str] = None) -> bytes:
if algorithm == 'lz4':
if TYPE_CHECKING:
from lz4 import frame
else:
lz4 = import_library('lz4', raise_error=True) # noqa: F841
from lz4 import frame
data = frame.compress(data)
elif algorithm == 'bz2':
import bz2
data = bz2.compress(data)
elif algorithm == 'lzma':
import lzma
data = lzma.compress(data)
elif algorithm == 'zlib':
import zlib
data = zlib.compress(data)
elif algorithm == 'gzip':
import gzip
data = gzip.compress(data)
return data
def _decompress_bytes(data: bytes, algorithm: Optional[str] = None) -> bytes:
if algorithm == 'lz4':
if TYPE_CHECKING:
from lz4 import frame
else:
lz4 = import_library('lz4', raise_error=True) # noqa: F841
from lz4 import frame
data = frame.decompress(data)
elif algorithm == 'bz2':
import bz2
data = bz2.decompress(data)
elif algorithm == 'lzma':
import lzma
data = lzma.decompress(data)
elif algorithm == 'zlib':
import zlib
data = zlib.decompress(data)
elif algorithm == 'gzip':
import gzip
data = gzip.decompress(data)
return data
def _get_compress_ctx(algorithm: Optional[str] = None) -> Optional[Callable]:
if algorithm == 'lz4':
if TYPE_CHECKING:
from lz4 import frame
else:
lz4 = import_library('lz4', raise_error=True) # noqa: F841
from lz4 import frame
def _fun(x: IO[bytes]):
return frame.LZ4FrameFile(x, 'wb')
compress_ctx = _fun
elif algorithm == 'gzip':
import gzip
def _fun(x: IO[bytes]):
return gzip.GzipFile(fileobj=x, mode='wb')
compress_ctx = _fun
elif algorithm == 'bz2':
import bz2
def _fun(x: IO[bytes]):
return bz2.BZ2File(filename=x, mode='wb')
compress_ctx = _fun
elif algorithm == 'lzma':
import lzma
def _fun(x: IO[bytes]):
return lzma.LZMAFile(filename=x, mode='wb')
compress_ctx = _fun
else:
compress_ctx = None
return compress_ctx