JuliaHEP · graeme-a-stewart · Mar 13, 2025 · Mar 12, 2025 · Mar 12, 2025 · Mar 12, 2025
diff --git a/Project.toml b/Project.toml
@@ -6,6 +6,7 @@ version = "0.4.5"
 [deps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
 CodecZlib = "944b1d66-785c-5afd-91f1-9de20f533193"
+CodecZstd = "6b39b394-51ab-5f42-8807-6242bab2b4c2"
 EnumX = "4e289a0a-7415-4d19-859d-a7e5c4648b56"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
@@ -25,7 +26,8 @@ JetVisualisation = "Makie"
 [compat]
 Accessors = "0.1.36"
 Aqua = "0.8"
-CodecZlib = "0.7.4"
+CodecZlib = "0.7"
+CodecZstd = "0.8"
 EDM4hep = "0.4.0"
 EnumX = "1.0.4"
 JSON = "0.21"
@@ -41,8 +43,8 @@ julia = "1.9"
 
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
 test = ["Aqua", "Test", "JSON"]
diff --git a/src/Utils.jl b/src/Utils.jl
@@ -1,6 +1,24 @@
 # Utility functions, which can be used by different top level scripts
 
 using CodecZlib
+using CodecZstd
+
+"""
+    open_with_stream(fname::AbstractString)
+
+Open a file with a stream decompressor if it is compressed with gzip or zstd,
+otherwise as a normal file.
+"""
+open_with_stream(fname::AbstractString) = begin
+    if endswith(fname, ".gz")
+        f = GzipDecompressorStream(open(fname))
+    elseif endswith(fname, ".zst")
+        f = ZstdDecompressorStream(open(fname))
+    else
+        f = open(fname)
+    end
+    f
+end
 
 """
     read_final_state_particles(fname; maxevents = -1, skipevents = 0, T=PseudoJet)
@@ -22,12 +40,7 @@ a `LorentzVector` type. Note, if T is not `PseudoJet`, the order of the
 arguments in the constructor must be `(t, x, y, z)`.
 """
 function read_final_state_particles(fname; maxevents = -1, skipevents = 0, T = PseudoJet)
-    f = open(fname)
-    if endswith(fname, ".gz")
-        @debug "Reading gzipped file $fname"
-        f = GzipDecompressorStream(f)
-    end
-
+    f = open_with_stream(fname)
     events = Vector{T}[]
 
     ipart = 1

diff --git a/test/_common.jl b/test/_common.jl
@@ -10,7 +10,6 @@ using Logging
 using LorentzVectorHEP
 using JSON
 using Test
-using CodecZlib
 
 logger = ConsoleLogger(stdout, Logging.Warn)
 global_logger(logger)
@@ -66,11 +65,7 @@ end
 
 """Read JSON file with fastjet jets in it"""
 function read_fastjet_outputs(fname)
-    f = open(fname)
-    if endswith(fname, ".gz")
-        @debug "Reading gzipped file $fname"
-        f = GzipDecompressorStream(f)
-    end
+    f = JetReconstruction.open_with_stream(fname)
     JSON.parse(f)
 end
 

diff --git a/test/data/file.txt b/test/data/file.txt
@@ -0,0 +1 @@
+Hello, JetReconstruction!
diff --git a/test/data/file.txt.gz b/test/data/file.txt.gz
diff --git a/test/data/file.txt.zst b/test/data/file.txt.zst
diff --git a/test/runtests.jl b/test/runtests.jl
@@ -29,6 +29,11 @@ function main()
     #
     # Now run a few tests with our examples
     # include("tests_examples.jl")
+
+    # Utility support tests
+    include("test-utils.jl")
+
+    # Substructure tests
     include("test-substructure.jl")
 
     # Test with Aqua (https://juliatesting.github.io/Aqua.jl/stable/)

diff --git a/test/test-utils.jl b/test/test-utils.jl
@@ -0,0 +1,14 @@
+# Tests of utility functions
+
+# Mainly just for "using"
+include("common.jl")
+
+const string_target = "Hello, JetReconstruction!"
+
+@testset "File streaming" begin
+    for fname in ["file.txt", "file.txt.gz", "file.txt.zst"]
+        fname = joinpath(@__DIR__, "data", fname)
+        string = readline(JetReconstruction.open_with_stream(fname))
+        @test string == string_target
+    end
+end