Move ZstdCompressor to a package extension

Project.toml

@@ -20,6 +20,12 @@ Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 URIs = "5c2747f8-b7ea-4ff2-ba2e-563bfd36b1d4"
 ZipArchives = "49080126-0e18-4c2a-b176-c102e4b3760c"
 
+[weakdeps]
+ChunkCodecLibZstd = "55437552-ac27-4d47-9aa3-63184e8fd398"
+
+[extensions]
+ChunkCodecLibZstdExt = "ChunkCodecLibZstd"
+
 [compat]
 AWSS3 = "0.10, 0.11"
 Blosc = "0.5, 0.6, 0.7"

ext/ChunkCodecLibZstdExt.jl

@@ -0,0 +1,60 @@
+#=
+# Zstd compression
+
+This file implements a Zstd compressor via ChunkCodecLibZstd.jl.
+
+=#
+module ChunkCodecLibZstdExt
+
+using ChunkCodecLibZstd: ZstdEncodeOptions, encode, decode, ChunkCodecCore
+using JSON: JSON
+import Zarr: Zarr, Compressor, getCompressor, zuncompress, zuncompress!, zcompress
+
+
+"""
+    ZstdCompressor(;level=0, checksum=false)
+Returns a `ZstdCompressor` struct that can serve as a Zarr array compressor. Keyword arguments are:
+* `level=0`: the compression level, regular levels are 1 to 22, 0 is a special value for default, there are also even faster negative levels.
+* `checksum=false`: flag to enable saving checksums.
+"""
+struct ZstdCompressor <: Compressor
+    config::ZstdEncodeOptions
+end
+
+ZstdCompressor(;level=0, checksum::Bool=false) = ZstdCompressor(ZstdEncodeOptions(;compressionLevel=level, checksum))
+
+function getCompressor(::Type{ZstdCompressor}, d::Dict)
+    ZstdCompressor(;
+        level=get(Returns(0), d, "level"),
+        checksum=Bool(get(Returns(false), d, "checksum")),
+    )
+end
+
+function zuncompress(a, z::ZstdCompressor, T)
+    result = decode(z.config.codec, a)
+    _reinterpret(Base.nonmissingtype(T),result)
+end
+
+function zuncompress!(data::DenseArray, compressed, z::ZstdCompressor)
+    dst = reinterpret(UInt8, vec(data))
+    n = length(dst)
+    n_decoded = something(ChunkCodecCore.try_decode!(z.config.codec, dst, compressed))::Int64
+    n_decoded == n || error("expected to decode $n bytes, only got $n_decoded bytes")
+    data
+end
+
+function zcompress(a, z::ZstdCompressor)
+    encode(z.config, reinterpret(UInt8, vec(a)))
+end
+
+JSON.lower(z::ZstdCompressor) = Dict("id"=>"zstd", "level" => z.config.compressionLevel, "checksum" => z.config.checksum)
+
+function register(::Type{ZstdCompressor})
+    Zarr.compressortypes["zstd"] = ZstdCompressor
+end
+
+function __init__()
+    register(ZstdCompressor)
+end
+
+end # module ChunkCodecLibZstdExt

src/Compressors/Compressors.jl

@@ -104,4 +104,4 @@ end
 
 JSON.lower(::NoCompressor) = nothing
 
-compressortypes[nothing] = NoCompressor
+compressortypes[nothing] = NoCompressor

src/Compressors/zstd.jl

@@ -1,49 +1,16 @@
-#=
-# Zstd compression
-
-This file implements a Zstd compressor via ChunkCodecLibZstd.jl.
-
-=#
-
-using ChunkCodecLibZstd: ZstdEncodeOptions, encode, decode, ChunkCodecCore
-
+# This is a stub for ChunkCodecLibZstdExt.ZstdCompressor()
 
 """
     ZstdCompressor(;level=0, checksum=false)
 Returns a `ZstdCompressor` struct that can serve as a Zarr array compressor. Keyword arguments are:
 * `level=0`: the compression level, regular levels are 1 to 22, 0 is a special value for default, there are also even faster negative levels.
 * `checksum=false`: flag to enable saving checksums.
 """
-struct ZstdCompressor <: Compressor
-    config::ZstdEncodeOptions
+function ZstdCompressor(; kwargs...)
+    ChunkCodecLibZstdExt = Base.get_extension(Zarr, :ChunkCodecLibZstdExt)
+    if isnothing(ChunkCodecLibZstdExt)
+        error("Please load ChunkCodecLibZstd.jl by typing `using ChunkCodecLibZstd`")
+    else
+        return ChunkCodecLibZstdExt.ZstdCompressor(; kwargs...)
+    end
 end
-
-ZstdCompressor(;level=0, checksum::Bool=false) = ZstdCompressor(ZstdEncodeOptions(;compressionLevel=level, checksum))
-
-function getCompressor(::Type{ZstdCompressor}, d::Dict)
-    ZstdCompressor(;
-        level=get(Returns(0), d, "level"),
-        checksum=Bool(get(Returns(false), d, "checksum")),
-    )
-end
-
-function zuncompress(a, z::ZstdCompressor, T)
-    result = decode(z.config.codec, a)
-    _reinterpret(Base.nonmissingtype(T),result)
-end
-
-function zuncompress!(data::DenseArray, compressed, z::ZstdCompressor)
-    dst = reinterpret(UInt8, vec(data))
-    n = length(dst)
-    n_decoded = something(ChunkCodecCore.try_decode!(z.config.codec, dst, compressed))::Int64
-    n_decoded == n || error("expected to decode $n bytes, only got $n_decoded bytes")
-    data
-end
-
-function zcompress(a, z::ZstdCompressor)
-    encode(z.config, reinterpret(UInt8, vec(a)))
-end
-
-JSON.lower(z::ZstdCompressor) = Dict("id"=>"zstd", "level" => z.config.compressionLevel, "checksum" => z.config.checksum)
-
-Zarr.compressortypes["zstd"] = ZstdCompressor

test/Project.toml

@@ -1,4 +1,5 @@
 [deps]
+ChunkCodecLibZstd = "55437552-ac27-4d47-9aa3-63184e8fd398"
 CondaPkg = "992eb4ea-22a4-4c89-a5bb-47a3300528ab"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 JSON = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"

test/python.jl

@@ -3,6 +3,9 @@
 ### We save some data in Julia and python and test if it is still the same
 ### when read from the other language 
 ###
+using Test
+using ChunkCodecLibZstd
+
 @testset "Python zarr implementation" begin
 
 import Mmap