From df7dd5de36a29d2acd238c8ae3e2292f31810ef9 Mon Sep 17 00:00:00 2001
From: Daniel Loos <dloos@bgc-jena.mpg.de>
Date: Thu, 29 Aug 2024 13:48:45 +0200
Subject: [PATCH] Fix collect of `ConcatDiskArray` (#427)

* Fix collect of concat arrays

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

diff --git a/src/Cubes/Cubes.jl b/src/Cubes/Cubes.jl
index 6666cde6..ff1cb239 100644
--- a/src/Cubes/Cubes.jl
+++ b/src/Cubes/Cubes.jl
@@ -169,7 +169,7 @@ function Base.propertynames(a::YAXArray, private::Bool=false)
     end
 end
 
-
+Base.Generator(f, A::YAXArray) = Base.Generator(f, parent(A))
 Base.ndims(a::YAXArray{<:Any,N}) where {N} = N
 Base.eltype(a::YAXArray{T}) where {T} = T
 function Base.permutedims(c::YAXArray, p) 
diff --git a/src/YAXTools.jl b/src/YAXTools.jl
index a0842dfd..e608ac17 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)
@@ -79,4 +79,4 @@ function Base.eltype(p::PickAxisArray{T}) where {T}
     end
 end
 Base.getindex(p::PickAxisArray, i::CartesianIndex) = p[i.I...]
-end
+end
\ 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