Coverage for src/evutils/io/_compression.py: 81%

1"""Transparent (de)compression helpers for path-based IO. 

2 

3Event files are frequently stored compressed (``foo.raw.zst``, 

4``events.csv.gz``). This module centralises the mapping from a compression 

5suffix to the right stdlib (or third-party) file wrapper so both the reader 

6(:mod:`evutils.io._source`) and the writer 

7(:class:`~evutils.io._event_writer.EventWriter`) can open them transparently. 

8 

9The *inner* extension still selects the event format: ``foo.raw.zst`` is an 

10EVT file read/written through a zstd stream. Use 

11:func:`strip_compression_suffix` to recover the inner name for format 

12detection. 

13 

14Supported suffixes: 

15 

16* ``.gz`` -- gzip (stdlib :mod:`gzip`) 

17* ``.xz`` -- lzma/xz (stdlib :mod:`lzma`) 

18* ``.bz2`` -- bzip2 (stdlib :mod:`bz2`) 

19* ``.zst`` -- zstandard: stdlib ``compression.zstd`` (Python 3.14+), falling 

20 back to the third-party ``zstandard`` or ``pyzstd`` packages if installed. 

21""" 

22from __future__ import annotations 

23 

24import io 

25from pathlib import Path 

26 

27#: Recognised compression suffixes (lower-case, incl. leading dot). 

28COMPRESSION_SUFFIXES = {".gz", ".zst", ".xz", ".bz2"} 

29 

30def is_compressed_path(path: "str | Path") -> bool: 

31 """Return True if ``path``'s final suffix is a known compression suffix.""" 

32 return Path(path).suffix.lower() in COMPRESSION_SUFFIXES 

33 

34def strip_compression_suffix(name: str) -> str: 

35 """Drop a trailing compression suffix: ``'foo.raw.zst'`` -> ``'foo.raw'``. 

36 

37 Names without a compression suffix are returned unchanged. 

38 """ 

39 p = Path(name) 

40 suffix = p.suffix 

41 if suffix.lower() in COMPRESSION_SUFFIXES: 

42 return name[:-len(suffix)] 

43 return name 

44 

45def _open_zstd(path: "str | Path", mode: str) -> "io.BufferedIOBase": 

46 """Open a ``.zst`` file, trying stdlib then third-party backends.""" 

47 try: 

48 from compression.zstd import ZstdFile # Python 3.14+ stdlib 

49 return ZstdFile(path, mode) 

50 except ImportError: 

51 pass 

52 try: 

53 import zstandard # third-party 

54 return zstandard.open(path, mode) 

55 except ImportError: 

56 pass 

57 try: 

58 import pyzstd # third-party 

59 return pyzstd.ZstdFile(path, mode) 

60 except ImportError: 

61 pass 

62 raise ImportError( 

63 "reading/writing '.zst' files requires zstd support: use Python 3.14+ " 

64 "(stdlib 'compression.zstd') or install the 'zstandard' or 'pyzstd' " 

65 "package" 

66 ) 

67 

68def open_compressed(path: "str | Path", mode: str = "rb") -> "io.BufferedIOBase": 

69 """Open a compressed file, dispatching on its suffix. 

70 

71 Parameters 

72 ---------- 

73 path 

74 Path whose final suffix is one of :data:`COMPRESSION_SUFFIXES`. 

75 mode 

76 Binary open mode (``'rb'`` / ``'wb'``). Text modes are not supported -- 

77 event codecs operate on bytes. 

78 

79 Returns 

80 ------- 

81 io.BufferedIOBase 

82 A binary, decompressing/compressing file object wrapping ``path``. 

83 

84 Raises 

85 ------ 

86 ValueError 

87 If the suffix is not a recognised compression suffix. 

88 ImportError 

89 If the backend for the suffix is unavailable (e.g. zstd). 

90 

91 """ 

92 suffix = Path(path).suffix.lower() 

93 if suffix == ".gz": 

94 import gzip 

95 return gzip.open(path, mode) 

96 if suffix == ".xz": 

97 import lzma 

98 return lzma.open(path, mode) 

99 if suffix == ".bz2": 

100 import bz2 

101 return bz2.open(path, mode) 

102 if suffix == ".zst": 

103 return _open_zstd(path, mode) 

104 raise ValueError( 

105 f"{suffix!r} is not a supported compression suffix " 

106 f"(expected one of {sorted(COMPRESSION_SUFFIXES)})" 

107 )