feat: Add zstd support for reading and writing (#95)

This uses the `compression.zstd` module in Python 3.14 and
`backports.zstd` (unconditionally installed from pypi) for lower
versions

Author: paulgessinger

pyproject.toml

@@ -24,7 +24,7 @@ classifiers = [
     "Topic :: Scientific/Engineering",
     "Intended Audience :: Developers",
 ]
-dependencies = ["numpy>=1.21", "packaging"]
+dependencies = ["numpy>=1.21", "packaging", "backports.zstd"]
 dynamic = ["version"]
 
 [project.urls]

src/pyhepmc/io.py

@@ -247,6 +247,14 @@ def __init__(
                 import lzma
 
                 open = lzma.open  # type:ignore
+            elif fn.endswith(".zst") or fn.endswith(".zstd"):
+                from sys import version_info
+
+                if version_info >= (3, 14):
+                    from compression import zstd
+                else:
+                    from backports import zstd
+                open = zstd.open
             else:
                 from builtins import open  # type:ignore
 

tests/test_io.py

@@ -16,6 +16,11 @@
 else:
     list_type = typing.List
 
+if version_info >= (3, 14):
+    from compression import zstd
+else:
+    from backports import zstd
+
 # this only does something if pyhepmc is compiled in debug mode
 hep.Setup.print_warnings = True
 
@@ -106,6 +111,23 @@ def test_pystream_5():
         pio.write(b"foo")
 
 
+def test_pystream_6(evt):
+    fn = "test_pystream_6.dat.zst"
+    with zstd.open(fn, "w") as f:
+        with pyiostream(f, 1000) as s:
+            with io.WriterAscii(s) as w:
+                w.write(evt)
+
+    with zstd.open(fn) as f:
+        with pyiostream(f, 1000) as s:
+            with io.ReaderAscii(s) as r:
+                evt2 = r.read()
+
+    assert evt == evt2
+
+    os.unlink(fn)
+
+
 def test_read_event_write_event(evt):
     oss = stringstream()
     with io.WriterAscii(oss) as f: