From 135506dabb8bdef423e85cf3eba5a90ea5cdc66b Mon Sep 17 00:00:00 2001
From: Alexander Barth <barth.alexander@gmail.com>
Date: Tue, 7 May 2024 12:07:21 +0200
Subject: [PATCH] add test of groups

---
 src/ZarrDatasets.jl |  1 +
 src/dataset.jl      | 40 ++++++++++++++++++-------
 src/variable.jl     |  3 +-
 test/runtests.jl    |  1 +
 test/test_groups.jl | 72 +++++++++++++++++++++++++++++++++++++++++++++
 5 files changed, 105 insertions(+), 12 deletions(-)
 create mode 100644 test/test_groups.jl

diff --git a/src/ZarrDatasets.jl b/src/ZarrDatasets.jl
index c7a3221..d13dd57 100644
--- a/src/ZarrDatasets.jl
+++ b/src/ZarrDatasets.jl
@@ -16,6 +16,7 @@ import CommonDataModel:
     defAttrib,
     defVar,
     defDim,
+    defGroup,
     dim,
     dimnames,
     iswritable,
diff --git a/src/dataset.jl b/src/dataset.jl
index 82a97ec..df436e0 100644
--- a/src/dataset.jl
+++ b/src/dataset.jl
@@ -18,14 +18,19 @@ CDM.dimnames(ds::ZarrDataset) = Tuple(String.(keys(ds.dimensions)))
 #     return ul
 # end
 
-# function _dim(ds::ZarrDataset,dimname::SymbolOrString)
-#     if haskey(ds.dimensions,name)
-#         return ds.dimensions[name]
-#     elseif ds.parentdataset !== nothing
-#         return _dim(ds.parentdataset,name)
-#     end
-#     return nothing
-# end
+function _dim(ds::ZarrDataset,dimname::SymbolOrString)
+    dimlen = get(ds.dimensions,Symbol(dimname),nothing)
+
+    if !isnothing(dimlen)
+        return dimlen
+    end
+
+    if ds.parentdataset !== nothing
+        return _dim(ds.parentdataset,dimname)
+    end
+
+    error("dimension $dimname is not defined")
+end
 
 CDM.dim(ds::ZarrDataset,dimname::SymbolOrString) = ds.dimensions[Symbol(dimname)]
 
@@ -50,8 +55,23 @@ function CDM.defAttrib(ds::ZarrDataset,name::SymbolOrString,value)
     storage[ds.zgroup.path,".zattrs"] = take!(io)
 end
 
+# groups
+
+function CDM.defGroup(ds::ZarrDataset,groupname::SymbolOrString; attrib = Dict())
+    _attrib = Dict{String,Any}(attrib)
+    zg = zgroup(ds.zgroup,String(groupname),attrs = _attrib)
+    dimensions = OrderedDict{Symbol,Int}()
+    return ZarrDataset(ds,zg,dimensions,ds.iswritable,ds.maskingvalue)
+end
+
 CDM.groupnames(ds::ZarrDataset) = keys(ds.zgroup.groups)
-CDM.group(ds::ZarrDataset,name::SymbolOrString) = ZarrDataset(ds.zgroup.groups,String(name),ds)
+
+function CDM.group(ds::ZarrDataset,groupname::SymbolOrString)
+    dimensions = OrderedDict{Symbol,Int}()
+    zg = ds.zgroup.groups[String(groupname)]
+    return ZarrDataset(ds,zg,dimensions,ds.iswritable,ds.maskingvalue)
+end
+
 
 
 CDM.parentdataset(ds::ZarrDataset) = ds.parentdataset
@@ -161,4 +181,4 @@ end
 export ZarrDataset
 export defDim
 export defVar
-#export defGroup
+export defGroup
diff --git a/src/variable.jl b/src/variable.jl
index fc89542..6eb6341 100644
--- a/src/variable.jl
+++ b/src/variable.jl
@@ -44,7 +44,6 @@ haschunks(v::ZarrVariable) = haschunks(v.zarray)
 eachchunk(v::CFVariable{T,N,<:ZarrVariable}) where {T,N} = eachchunk(v.var)
 haschunks(v::CFVariable{T,N,<:ZarrVariable}) where {T,N} = haschunks(v.var)
 
-
 function CDM.defVar(ds::ZarrDataset,name::SymbolOrString,vtype::DataType,dimensionnames; chunksizes=nothing, attrib = Dict(), fillvalue = nothing, kwargs...)
     @assert iswritable(ds)
 
@@ -56,7 +55,7 @@ function CDM.defVar(ds::ZarrDataset,name::SymbolOrString,vtype::DataType,dimensi
     _attrib["_ARRAY_DIMENSIONS"] = reverse(dimensionnames)
 
     _size = ntuple(length(dimensionnames)) do i
-        ds.dimensions[Symbol(dimensionnames[i])]
+        _dim(ds,dimensionnames[i])
     end
 
     if isnothing(chunksizes)
diff --git a/test/runtests.jl b/test/runtests.jl
index 63a6493..3b6cfb6 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -5,5 +5,6 @@ using ZarrDatasets
     include("test_cdm.jl")
     include("test_multifile.jl")
     include("test_write.jl")
+    include("test_groups.jl")
     include("test_aqua.jl")
 end
diff --git a/test/test_groups.jl b/test/test_groups.jl
new file mode 100644
index 0000000..99f8e91
--- /dev/null
+++ b/test/test_groups.jl
@@ -0,0 +1,72 @@
+using ZarrDatasets
+using Test
+
+data = rand(Int32,3,5)
+data2 = rand(Int8,3,5)
+data3 = rand(Int16,4)
+
+fname = tempname()
+mkdir(fname)
+
+
+TDS = ZarrDataset
+
+#using NCDatasets
+#TDS = NCDataset
+ds = TDS(fname,"c")
+defDim(ds,"lon",3)
+defDim(ds,"lat",5)
+
+attrib = Dict(
+    "units" => "m/s",
+    "long_name" => "test",
+)
+
+varname = "var2"
+dimensionnames = ("lon","lat")
+vtype = Int32
+
+zv = defVar(ds,varname,vtype,dimensionnames, attrib = attrib)
+zv[:,:] = data
+zv.attrib["number"] = 12
+ds.attrib["history"] = "test"
+
+
+group_name = "sub-group"
+attrib = Dict("int_attrib" => 42)
+
+dsg = defGroup(ds,group_name; attrib = attrib)
+
+defDim(dsg,"time",length(data3))
+
+zvg = defVar(dsg,"data2",eltype(data2),("lon","lat"))
+zvg[:,:] = data2
+zvg.attrib["standard_name"] = "test"
+
+zv3 = defVar(dsg,"data3",eltype(data3),("time",))
+zv3[:] = data3
+
+
+@test_throws Exception defVar(dsg,"data4",Int8,("dimension_does_not_exists",))
+
+
+io = IOBuffer()
+show(io,ds)
+s = String(take!(io))
+@test occursin("sub-group",s)
+
+close(ds)
+
+# load data from group
+
+ds = TDS(fname,"r")
+
+@test ds["var2"][:,:] == data
+@test ds.group[group_name]["data2"][:,:] == data2
+@test ds.group[group_name]["data2"].attrib["standard_name"] == "test"
+@test ds.group[group_name]["data3"][:] == data3
+
+@test ds.group[group_name].attrib["int_attrib"] == 42
+
+display(ds)
+close(ds)