47 lines
1.2 KiB
Python
47 lines
1.2 KiB
Python
"""
|
|
Various helpers for compression
|
|
"""
|
|
from pathlib import Path, PosixPath
|
|
from typing import Union
|
|
|
|
PathIsh = Union[Path, str]
|
|
|
|
|
|
def _zstd_open(path: Path):
|
|
import zstandard as zstd # type: ignore
|
|
fh = path.open('rb')
|
|
dctx = zstd.ZstdDecompressor()
|
|
reader = dctx.stream_reader(fh)
|
|
return reader
|
|
|
|
|
|
def kopen(path: PathIsh, *args, **kwargs): # TODO is it bytes stream??
|
|
pp = Path(path)
|
|
suf = pp.suffix
|
|
if suf in {'.xz'}:
|
|
import lzma
|
|
return lzma.open(pp, *args, **kwargs)
|
|
elif suf in {'.zip'}:
|
|
from zipfile import ZipFile
|
|
return ZipFile(pp).open(*args, **kwargs)
|
|
elif suf in {'.lz4'}:
|
|
import lz4.frame # type: ignore
|
|
return lz4.frame.open(str(pp))
|
|
elif suf in {'.zstd'}:
|
|
return _zstd_open(pp)
|
|
else:
|
|
return pp.open(*args, **kwargs)
|
|
|
|
|
|
class CPath(PosixPath):
|
|
"""
|
|
Ugh. So, can't override Path because of some _flavour thing.
|
|
Path only has _accessor and _closed slots, so can't directly set .open method
|
|
_accessor.open has to return file descriptor, doesn't work for compressed stuff.
|
|
"""
|
|
def open(self, *args, **kwargs):
|
|
# TODO assert read only?
|
|
return kopen(str(self))
|
|
|
|
|
|
open = kopen # TODO FIXME remove?
|