Bugfix and test _sourcetrait (#870)

* bugfix and test sourcetrait better * actually run sources tests * more tests
rafaqz · Jan 23, 2025 · d4d42e5 · d4d42e5
1 parent b55b28f
commit d4d42e5
Show file tree

Hide file tree

Showing 3 changed files with 91 additions and 4 deletions.
diff --git a/src/sources/sources.jl b/src/sources/sources.jl
@@ -65,12 +65,27 @@ function Base.showerror(io::IO, e::BackendException)
 end
 
 # Get the source backend for a file extension, falling back to GDALsource
+_sourcetrait(filename::AbstractString, s::Symbol) = _sourcetrait(s)
 _sourcetrait(filename::AbstractString, s::Source) = s
-_sourcetrait(filename::AbstractString, s) = _sourcetrait(s)
+_sourcetrait(filename::AbstractString, ::Type{S}) where S<:Source = S()
 _sourcetrait(filename::AbstractString, ::Union{Nothing,NoKW}) = _sourcetrait(filename)
-_sourcetrait(filename::AbstractString) = get(EXT2SOURCE, splitext(filename)[2], GDALsource())
+_sourcetrait(filename::AbstractString, ext::AbstractString) = get(EXT2SOURCE, ext, GDALsource())
+function _sourcetrait(filename::AbstractString)
+    default = GDALsource()
+    stem, ext = splitext(filename)
+    str = if ext == ""  
+        # Handle e.g. "x.zarr/" directories
+        if isdirpath(stem) 
+            return _sourcetrait(dirname(stem))
+        else
+            stem
+        end
+    else
+        ext
+    end
+    return get(EXT2SOURCE, str, default)
+end
 _sourcetrait(filenames::NamedTuple) = _sourcetrait(first(filenames))
-_sourcetrait(filename, ext) = get(EXT2SOURCE, ext, GDALsource())
 _sourcetrait(source::Source) = source
 _sourcetrait(source::Type{<:Source}) = source()
 function _sourcetrait(name::Symbol) 

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -26,7 +26,7 @@ end
 @time @safetestset "warp" begin include("warp.jl") end
 @time @safetestset "cellarea" begin include("cellarea.jl") end
 
-# CommondataModel sources
+@time @safetestset "sources" begin include("sources/sources.jl") end
 @time @safetestset "commondatamodel" begin include("sources/commondatamodel.jl") end
 @time @safetestset "ncdatasets" begin include("sources/ncdatasets.jl") end
 @time @safetestset "zarr" begin include("sources/zarr.jl") end

diff --git a/test/sources/sources.jl b/test/sources/sources.jl
@@ -0,0 +1,72 @@
+using Rasters, Test
+using Rasters: _sourcetrait
+
+@test _sourcetrait("x", ".nc") == Rasters.NCDsource()
+@test _sourcetrait("x", ".nc4") == Rasters.NCDsource()
+@test _sourcetrait("x", ".h5") == Rasters.NCDsource()
+@test _sourcetrait("x", ".grd") == Rasters.GRDsource()
+@test _sourcetrait("x", ".gri") == Rasters.GRDsource()
+@test _sourcetrait("x", ".grib") == Rasters.GRIBsource()
+@test _sourcetrait("x", ".tif") == Rasters.GDALsource()
+@test _sourcetrait("x", ".zarr") == Rasters.Zarrsource()
+
+@test _sourcetrait("x", :netcdf) == Rasters.NCDsource()
+@test _sourcetrait("x", :grd) == Rasters.GRDsource()
+@test _sourcetrait("x", :grib) == Rasters.GRIBsource()
+@test _sourcetrait("x", :gdal) == Rasters.GDALsource()
+@test _sourcetrait("x", :zarr) == Rasters.Zarrsource()
+
+@test _sourcetrait("x", Rasters.NCDsource()) == Rasters.NCDsource()
+@test _sourcetrait("x", Rasters.GRDsource()) == Rasters.GRDsource()
+@test _sourcetrait("x", Rasters.GRIBsource()) == Rasters.GRIBsource()
+@test _sourcetrait("x", Rasters.GDALsource()) == Rasters.GDALsource()
+@test _sourcetrait("x", Rasters.Zarrsource()) == Rasters.Zarrsource()
+
+@test _sourcetrait("x", Rasters.NCDsource) == Rasters.NCDsource()
+@test _sourcetrait("x", Rasters.GRDsource) == Rasters.GRDsource()
+@test _sourcetrait("x", Rasters.GRIBsource) == Rasters.GRIBsource()
+@test _sourcetrait("x", Rasters.GDALsource) == Rasters.GDALsource()
+@test _sourcetrait("x", Rasters.Zarrsource) == Rasters.Zarrsource()
+
+@test _sourcetrait("x", :netcdf) == Rasters.NCDsource()
+@test _sourcetrait("x", :grd) == Rasters.GRDsource()
+@test _sourcetrait("x", :grib) == Rasters.GRIBsource()
+@test _sourcetrait("x", :gdal) == Rasters.GDALsource()
+@test _sourcetrait("x", :zarr) == Rasters.Zarrsource()
+
+@test _sourcetrait(".nc") == Rasters.NCDsource()
+@test _sourcetrait(".nc4") == Rasters.NCDsource()
+@test _sourcetrait(".h5") == Rasters.NCDsource()
+@test _sourcetrait(".grd") == Rasters.GRDsource()
+@test _sourcetrait(".grib") == Rasters.GRIBsource()
+@test _sourcetrait(".tif") == Rasters.GDALsource()
+@test _sourcetrait(".zarr") == Rasters.Zarrsource()
+@test _sourcetrait(".zarr/") == Rasters.Zarrsource()
+
+@test _sourcetrait("x.nc") == Rasters.NCDsource()
+@test _sourcetrait("x.nc4") == Rasters.NCDsource()
+@test _sourcetrait("x.h5") == Rasters.NCDsource()
+@test _sourcetrait("x.grd") == Rasters.GRDsource()
+@test _sourcetrait("x.gri") == Rasters.GRDsource()
+@test _sourcetrait("x.grib") == Rasters.GRIBsource()
+@test _sourcetrait("x.tif") == Rasters.GDALsource()
+@test _sourcetrait("x.zarr") == Rasters.Zarrsource()
+@test _sourcetrait("x.zarr/") == Rasters.Zarrsource()
+
+@test _sourcetrait(:netcdf) == Rasters.NCDsource()
+@test _sourcetrait(:grd) == Rasters.GRDsource()
+@test _sourcetrait(:grib) == Rasters.GRIBsource()
+@test _sourcetrait(:gdal) == Rasters.GDALsource()
+@test _sourcetrait(:zarr) == Rasters.Zarrsource()
+
+@test _sourcetrait(Rasters.NCDsource()) == Rasters.NCDsource()
+@test _sourcetrait(Rasters.GRDsource()) == Rasters.GRDsource()
+@test _sourcetrait(Rasters.GRIBsource()) == Rasters.GRIBsource()
+@test _sourcetrait(Rasters.GDALsource()) == Rasters.GDALsource()
+@test _sourcetrait(Rasters.Zarrsource()) == Rasters.Zarrsource()
+
+@test _sourcetrait(Rasters.NCDsource) == Rasters.NCDsource()
+@test _sourcetrait(Rasters.GRDsource) == Rasters.GRDsource()
+@test _sourcetrait(Rasters.GRIBsource) == Rasters.GRIBsource()
+@test _sourcetrait(Rasters.GDALsource) == Rasters.GDALsource()
+@test _sourcetrait(Rasters.Zarrsource) == Rasters.Zarrsource()