diff --git a/src/openstack.jl b/src/openstack.jl index 79bb5be32..34c3474fd 100644 --- a/src/openstack.jl +++ b/src/openstack.jl @@ -31,13 +31,19 @@ DD.data_eltype(::OpenStack{<:Any,<:Any,T}) where T = T Base.keys(::OpenStack{<:Any,K}) where K = K # TODO test this, and does it make sense to return an iterator here? -Base.values(os::OpenStack{<:Any,K}) where K = (os[k] for k in K) +Base.values(os::OpenStack{<:Any,K}) where K = map(k -> os[k], K) +Base.NamedTuple(os::OpenStack{<:Any,K}) where K = NamedTuple{K}(values(os)) +Base.map(f, os::OpenStack) = map(f, NamedTuple(os)) # Indexing OpenStack returns memory-backed Raster. +Base.iterate(os::OpenStack) = iterate(os, 1) +Base.iterate(os::OpenStack, i::Integer) = i > length(os) ? nothing : (os[i], i + 1) +Base.length(os::OpenStack{<:Any,K}) where K = length(K) function Base.getindex(os::OpenStack{<:Any,K}, key::Symbol) where K mods = os.mods - if mods isa AbstractModifications - _maybe_modify(dataset(os)[key], mods) + var = _open(os, key) + return if mods isa AbstractModifications + _maybe_modify(var, mods) else - _maybe_modify(dataset(os)[key], NamedTuple{K}(mods)[key]) + _maybe_modify(var, NamedTuple{K}(mods)[key]) end end diff --git a/src/sources/commondatamodel.jl b/src/sources/commondatamodel.jl index 8236f544e..d3bdebf3f 100644 --- a/src/sources/commondatamodel.jl +++ b/src/sources/commondatamodel.jl @@ -83,12 +83,16 @@ function _open(f, source::CDMsource, ds::AbstractDataset; cleanreturn(f(g)) else key = string(_name_or_firstname(g, name)) - v = CDM.variable(g, key) + v = _open(g, key) _open(f, source, v; mod) end end _open(f, ::CDMsource, var::AbstractArray; mod=NoMod(), kw...) = cleanreturn(f(_maybe_modify(var, mod))) +_open(os::OpenStack{<:CDMsource}, key::Union{AbstractString,Symbol}) = _open(dataset(os), key) +# Both needed for ambiguity +_open(os::AbstractDataset, key::AbstractString) = CDM.variable(os, key) +_open(os::AbstractDataset, key::Symbol) = CDM.variable(os, key) # This allows arbitrary group nesting _getgroup(ds, ::Union{Nothing,NoKW}) = ds @@ -129,7 +133,7 @@ end function _layers(ds::AbstractDataset, ::NoKW=nokw, ::NoKW=nokw) nondim = _nondimnames(ds) grid_mapping = String[] - vars = map(k -> CDM.variable(ds, k), nondim) + vars = map(k -> _open(ds, k), nondim) attrs = map(CDM.attribs, vars) for attr in attrs if haskey(attr, "grid_mapping") @@ -146,7 +150,7 @@ end _layers(ds::AbstractDataset, names, ::NoKW) = _layers(ds, collect(names), nokw) function _layers(ds::AbstractDataset, names::Vector, ::NoKW) - vars = map(n -> CDM.variable(ds, n), names) + vars = map(n -> _open(ds, n), names) attrs = map(CDM.attribs, vars) (; names, vars, attrs) end diff --git a/src/stack.jl b/src/stack.jl index d918315a0..d17c8aa3d 100644 --- a/src/stack.jl +++ b/src/stack.jl @@ -542,11 +542,7 @@ end # Open a single file stack function Base.open(f::Function, st::AbstractRasterStack{K}; kw...) where K ost = OpenStack(parent(st)) - # TODO is this needed? - layers = map(K) do k - ost[k] - end |> NamedTuple{K} - out = f(rebuild(st; data=layers)) + out = f(rebuild(st; data=ost)) close(ost) return out end diff --git a/test/sources/ncdatasets.jl b/test/sources/ncdatasets.jl index 6a405611e..f97210e5d 100644 --- a/test/sources/ncdatasets.jl +++ b/test/sources/ncdatasets.jl @@ -1,4 +1,4 @@ -using Rasters, DimensionalData, Test, Statistics, Dates, CFTime, Plots, CommonDataModel +using Rasters, DimensionalData, Test, Statistics, Dates, CFTime, Plots, CommonDataModel, DiskArrays using Rasters.Lookups, Rasters.Dimensions using Rasters.DiskArrays @@ -626,4 +626,21 @@ if !haskey(ENV, "CI") && isfile(h5path) end end +@testset "OpenStack indexing" begin + DiskArrays.allowscalar(true) + filename = tempname() * ".nc" + rs = Rasters.create( + filename, (;a = Float64), (X(1:10),); + fill=NaN, + missingval=NaN + ) + rs_lazy = RasterStack(filename; lazy=true, raw=true) + ras_lazy = Raster(filename; lazy=true, raw=true) + + @test isequal(open(first, rs), open(first, rs_lazy)) + @test isequal(open(first, rs_lazy), (; a=NaN)) + @test isnan(open(first, ras_lazy)) + DiskArrays.allowscalar(false) +end + nothing