From 782146e144bc44a1cf792c4e49c2d0fac0b78935 Mon Sep 17 00:00:00 2001
From: Daniel Loos <dloos@bgc-jena.mpg.de>
Date: Mon, 26 Aug 2024 12:02:12 +0200
Subject: [PATCH] Fix collect of concat arrays

---
 src/YAXTools.jl                |  4 +++-
 test/Cubes/transformedcubes.jl | 19 ++++++++++++++++++-
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/src/YAXTools.jl b/src/YAXTools.jl
index a0842dfd..747c36f0 100644
--- a/src/YAXTools.jl
+++ b/src/YAXTools.jl
@@ -14,7 +14,7 @@ struct Window
 end
 
 
-function PickAxisArray(parent, indmask, perm = nothing)
+function PickAxisArray(parent, indmask, perm=nothing)
     f = findall(isequal(true), indmask)
     f2 = findall(isequal(Colon()), indmask)
     f3 = findall(i -> isa(i, Tuple{Int,Int}), indmask)
@@ -80,3 +80,5 @@ function Base.eltype(p::PickAxisArray{T}) where {T}
 end
 Base.getindex(p::PickAxisArray, i::CartesianIndex) = p[i.I...]
 end
+
+Base.Generator(f, A::YAXArray) = Base.Generator(f, parent(A))
\ No newline at end of file
diff --git a/test/Cubes/transformedcubes.jl b/test/Cubes/transformedcubes.jl
index c503519b..b0023e95 100644
--- a/test/Cubes/transformedcubes.jl
+++ b/test/Cubes/transformedcubes.jl
@@ -23,7 +23,24 @@
         @test caxes(newcube) == (X(1.0:4.0), Y([1, 2, 3, 4, 5]), Z(["A", "B", "C"]))
         @test ndims(newcube) == 3
         @test size(newcube) == (4, 5, 3)
-        @test newcube[:, :, :] == cat(data..., dims = 3)
+        @test newcube[:, :, :] == cat(data..., dims=3)
         @test getattributes(newcube) == reduce(merge, props)
     end
+    @testset "ConcatDiskArray" begin
+        using Zarr
+        lon_range = X(-180:180)
+        lat_range = Y(-90:90)
+        data = [exp(cosd(lon)) + 3 * (lat / 90) for lon in lon_range, lat in lat_range]
+        a = YAXArray((lon_range, lat_range), data)
+        ds_ram = Dataset(; properties=Dict(), a)
+        path = tempname()
+        savedataset(ds_ram; path=path)
+        ds_disk = open_dataset(path)
+        a_ram = cat(ds_ram.a[X=1:100], ds_ram.a[X=101:200], dims=:X)
+        a_disk = cat(ds_disk.a[X=1:100], ds_disk.a[X=101:200], dims=:X)
+
+        @test collect(x for x in a_disk) == collect(x for x in a_ram)
+
+        rm(path, recursive=true)
+    end
 end