Coverage for tests/io/test_compression.py: 99%
135 statements
« prev ^ index » next coverage.py v7.15.1, created at 2026-07-18 05:24 +0000
1"""Transparent compression: reading and writing compressed event files.
3Covers auto-open by path (``.gz`` / ``.zst`` / ``.xz`` / ``.bz2``), passing an
4already-open compressed file object as source, CSV over a compressed stream,
5and seeking over a compressed (non-seekable) native source.
6"""
7import gzip
8import io
10import numpy as np
11import pytest
13from typing import Any
14from evutils.io import EventReader, EventWriter
15from evutils.io._compression import (
16 COMPRESSION_SUFFIXES,
17 is_compressed_path,
18 open_compressed,
19 strip_compression_suffix,
20)
21from evutils.types import Event_dtype
24def _has_zstd() -> bool:
25 try:
26 open_compressed("x.zst", "rb") # raises ImportError if no backend
27 except ImportError:
28 return False
29 except Exception:
30 # FileNotFoundError etc. means a backend *is* available.
31 return True
32 return True
35# Suffixes always available via stdlib, plus zst when a backend exists.
36_SUFFIXES = [".gz", ".xz", ".bz2"]
37if _has_zstd():
38 _SUFFIXES.append(".zst")
41def make_events(n: int = 5000, t_max: int = 100_000, seed: int = 42) -> Any:
42 rng = np.random.default_rng(seed)
43 ev = np.zeros(n, dtype=Event_dtype)
44 ev["t"] = np.sort(rng.integers(0, t_max, n))
45 ev["x"] = rng.integers(0, 1280, n)
46 ev["y"] = rng.integers(0, 720, n)
47 ev["p"] = rng.integers(0, 2, n)
48 return ev
51def assert_events_equal(out: Any, ref: Any, context: str = "") -> None:
52 out = np.asarray(out)
53 assert len(out) == len(ref), f"{context}: length {len(out)} != {len(ref)}"
54 for f in ("t", "x", "y", "p"):
55 assert np.array_equal(out[f], ref[f]), f"{context}: field {f!r} differs"
58# --------------------------------------------------------------------------- #
59# Helper unit tests
60# --------------------------------------------------------------------------- #
61def test_is_compressed_path() -> None:
62 assert is_compressed_path("foo.raw.zst")
63 assert is_compressed_path("foo.csv.gz")
64 assert is_compressed_path("a.xz")
65 assert is_compressed_path("a.bz2")
66 assert not is_compressed_path("foo.raw")
67 assert not is_compressed_path("foo.csv")
68 # exactly the documented set
69 assert COMPRESSION_SUFFIXES == {".gz", ".zst", ".xz", ".bz2"}
72def test_strip_compression_suffix() -> None:
73 assert strip_compression_suffix("foo.raw.zst") == "foo.raw"
74 assert strip_compression_suffix("events.csv.gz") == "events.csv"
75 assert strip_compression_suffix("bar.raw") == "bar.raw" # unchanged
76 assert strip_compression_suffix("nested/foo.dat.xz") == "nested/foo.dat"
79def test_open_compressed_rejects_unknown_suffix() -> None:
80 with pytest.raises(ValueError):
81 open_compressed("foo.raw", "rb")
84# --------------------------------------------------------------------------- #
85# Round-trip: write compressed by path, read back by path
86# --------------------------------------------------------------------------- #
87@pytest.mark.parametrize("suffix", _SUFFIXES)
88def test_evt_roundtrip_compressed(tmp_path: Any, suffix: str) -> None:
89 ev = make_events()
90 p = tmp_path / f"out.raw{suffix}"
91 with EventWriter(p, format="evt3") as w:
92 w.write(ev)
93 with EventReader(p) as r:
94 out = np.asarray(r.read_all())
95 assert_events_equal(out, ev, suffix)
96 # first + last events match, as required.
97 assert out["t"][0] == ev["t"][0] and out["t"][-1] == ev["t"][-1]
100@pytest.mark.parametrize("suffix", _SUFFIXES)
101def test_compressed_matches_uncompressed(tmp_path: Any, suffix: str) -> None:
102 """A compressed read yields exactly the same events as the plain read."""
103 ev = make_events()
104 plain = tmp_path / "out.raw"
105 comp = tmp_path / f"out.raw{suffix}"
106 with EventWriter(plain, format="evt3") as w:
107 w.write(ev)
108 with EventWriter(comp, format="evt3") as w:
109 w.write(ev)
110 with EventReader(plain) as r:
111 a = np.asarray(r.read_all()).copy()
112 with EventReader(comp) as r:
113 b = np.asarray(r.read_all())
114 assert_events_equal(b, a, suffix)
117# --------------------------------------------------------------------------- #
118# Passing an already-open compressed file object as source
119# --------------------------------------------------------------------------- #
120def test_read_open_gzip_fileobj(tmp_path: Any) -> None:
121 ev = make_events()
122 p = tmp_path / "out.raw.gz"
123 with EventWriter(p, format="evt3") as w:
124 w.write(ev)
125 with gzip.open(p, "rb") as f:
126 with EventReader(f) as r:
127 out = np.asarray(r.read_all())
128 assert_events_equal(out, ev, "gzip-fileobj")
131def test_read_bytesio_of_gzip(tmp_path: Any) -> None:
132 """A GzipFile wrapping an in-memory buffer (no name) sniffs by content."""
133 ev = make_events()
134 p = tmp_path / "out.raw.gz"
135 with EventWriter(p, format="evt3") as w:
136 w.write(ev)
137 data = p.read_bytes()
138 gz = gzip.GzipFile(fileobj=io.BytesIO(data)) # no .name -> content sniff
139 with EventReader(gz) as r:
140 out = np.asarray(r.read_all())
141 assert_events_equal(out, ev, "gzip-bytesio")
144# --------------------------------------------------------------------------- #
145# CSV over a compressed (non-seekable) stream: header + no-header
146# --------------------------------------------------------------------------- #
147def test_csv_gzip_with_header(tmp_path: Any) -> None:
148 ev = make_events()
149 p = tmp_path / "events.csv.gz"
150 with EventWriter(p) as w: # CSV encoder writes a header by default
151 w.write(ev)
152 with EventReader(p) as r:
153 out = np.asarray(r.read_all())
154 assert_events_equal(out, ev, "csv.gz+header")
157def test_csv_gzip_no_header(tmp_path: Any) -> None:
158 ev = make_events()
159 p = tmp_path / "events.csv.gz"
160 with EventWriter(p, header=False) as w:
161 w.write(ev)
162 with EventReader(p) as r:
163 out = np.asarray(r.read_all())
164 assert_events_equal(out, ev, "csv.gz-noheader")
167def test_csv_open_gzip_fileobj(tmp_path: Any) -> None:
168 ev = make_events()
169 p = tmp_path / "events.csv.gz"
170 with EventWriter(p) as w:
171 w.write(ev)
172 with gzip.open(p, "rb") as f:
173 with EventReader(f) as r:
174 out = np.asarray(r.read_all())
175 assert_events_equal(out, ev, "csv gzip-fileobj")
178# --------------------------------------------------------------------------- #
179# Seeking over a compressed (non-seekable) native source
180# --------------------------------------------------------------------------- #
181def test_seek_over_compressed_evt(tmp_path: Any) -> None:
182 ev = make_events()
183 p = tmp_path / "out.raw.gz"
184 with EventWriter(p, format="evt3") as w:
185 w.write(ev)
186 with EventReader(p) as r:
187 landed = r.seek(n=1000)
188 chunk = np.asarray(r.read(n_events=5))
189 assert landed.ts == ev["t"][1000]
190 assert chunk["t"][0] == ev["t"][1000]
193# --------------------------------------------------------------------------- #
194# Chunked write over a compressed stream
195# --------------------------------------------------------------------------- #
196def test_chunked_write_compressed(tmp_path: Any) -> None:
197 ev = make_events()
198 p = tmp_path / "out.raw.bz2"
199 with EventWriter(p, format="evt3") as w:
200 for part in np.array_split(ev, 11):
201 w.write(part)
202 with EventReader(p) as r:
203 out = np.asarray(r.read_all())
204 assert_events_equal(out, ev, "chunked.bz2")