Carry dimension types through the axis interface: dimtype and a 6-arg yaxcreate

ext/DimensionalDataExt.jl

@@ -1,13 +1,16 @@
 module DimensionalDataExt
 using DimensionalData: DimArray, DimensionalData, data, Dim, metadata
-import YAXArrayBase: dimname, dimnames, dimvals, iscontdim, getattributes, getdata, yaxcreate
+import YAXArrayBase: dimname, dimnames, dimvals, dimtype, iscontdim, getattributes, getdata, yaxcreate
 _dname(::DimensionalData.Dim{N}) where N = N
 _dname(d::DimensionalData.Dimension) = DimensionalData.name(d)
 dimname(x::DimArray, i) = _dname(DimensionalData.dims(x)[i])
 
 
 dimvals(x::DimArray,i) = DimensionalData.dims(x)[i].val
 
+# the dimension's type with its parameters stripped: `D1` for a `D1`, `Dim{:lon}` for a `Dim{:lon}`
+dimtype(x::DimArray, i) = DimensionalData.basetypeof(DimensionalData.dims(x)[i])
+
 getdata(x::DimArray) = data(x)
 
 getattributes(x::DimArray) = metadata(x)
@@ -18,4 +21,13 @@ function yaxcreate(::Type{<:DimArray},data,dnames,dvals,atts)
   end
   DimArray(data,d,metadata = atts)
 end
+
+# with dimension types from the source where it has them
+function yaxcreate(::Type{<:DimArray},data,dnames,dtypes,dvals,atts)
+  d = ntuple(ndims(data)) do i
+    T = dtypes[i]
+    T === nothing ? Dim{Symbol(dnames[i])}(dvals[i]) : T(dvals[i])
+  end
+  DimArray(data,d,metadata = atts)
+end
 end

src/YAXArrayBase.jl

@@ -26,7 +26,7 @@ end
 
 
 export dimvals, dimname, dimnames, iscontdim, iscompressed,
-getattributes, getdata, yaxconvert
+getattributes, getdata, yaxconvert, dimtype
 export get_var_handle, get_varnames, get_var_dims, get_var_attrs,
 create_empty, add_var, allow_parallel_write,
 to_dataset, allow_missings

src/axisarrays/axisinterface.jl

@@ -19,6 +19,16 @@ dimname(::AbstractArray,i) = Symbol("Dim_",i)
 #Mandatory interface ends here
 
 # Optional methods
+"""
+    dimtype(x, i)
+
+Returns the dimension type of the i-th dimension of the array `x`, for targets that model
+dimensions as types (DimensionalData's `X` and `Ti`, or a user's `@dim D1`), or `nothing`
+when the array has only a name for it, in which case such targets use `Dim{dimname(x, i)}`.
+The fallback returns `nothing`.
+"""
+dimtype(x, i) = nothing
+
 """
     iscontdim(x, i)
 
@@ -78,12 +88,18 @@ Converts an AbstractArray x that implements the interface to type T.
 """
 function yaxconvert(T::Type{<:Any},x)
   data = getdata(x)
-  yaxcreate(T, data, dimnames(x),dimvals.(Ref(x),1:ndims(data)),getattributes(x))
+  n = ndims(data)
+  yaxcreate(T, data, dimnames(x), ntuple(i -> dimtype(x, i), n), dimvals.(Ref(x), 1:n), getattributes(x))
 end
 
 """
     yaxcreate(T::Type,data,dnames,dvals,attributes)
+    yaxcreate(T::Type,data,dnames,dtypes,dvals,attributes)
 
-Creates a new array with the given dimension names, values and attributes.
+Creates a new array with the given dimension names, values and attributes. In the
+six-argument form `dtypes[i]` is `dimtype(x, i)` of the source: a dimension type to
+construct the axis with, or `nothing`. Targets that do not model dimension types need only
+implement the five-argument form; the six-argument default forwards to it.
 """
 function yaxcreate(T,data,dname,dvals,attributes) end
+yaxcreate(T, data, dnames, dtypes, dvals, attributes) = yaxcreate(T, data, dnames, dvals, attributes)

test/arrays.jl

@@ -53,6 +53,15 @@ end
     @test dimvals(d,1) == 0.5:1.0:2.5
     @test dimvals(d,2) == 1.5:0.5:3.0
     @test getattributes(d) == Dict{String,Any}("a1"=>5, "a2"=>"att")
+    # a source without dimension types gives `nothing`, and the target then uses Dim{name}
+    @test dimtype(M(), 1) === nothing
+    @test typeof(dims(d, 1)) <: Dim{:x}
+    # a DimArray source keeps its dimension types through the conversion pipeline
+    DimensionalData.@dim D1
+    dd = DimArray(rand(2, 3), (D1(["a", "b"]), X(1:3)); metadata = Dict{String,Any}())
+    @test dimtype(dd, 1) === D1 && dimtype(dd, 2) === X
+    back = yaxconvert(DimArray, dd)
+    @test typeof(dims(back)) == typeof(dims(dd)) && back == dd
 end
 
 @testitem "ArchGDAL" begin