Custom types for wload to be used with produce_or_load

src/DrWatson.jl

@@ -35,8 +35,21 @@ function wsave(filename, data...; kwargs...)
     return _wsave(filename, data...; kwargs...)
 end
 
-"Currently equivalent with `FileIO.load`."
+"""
+    wload([::DataType], args...; kwargs...)
+
+Default fallback is `FileIO.load`, discarding the (optional) first argument.
+Extend this for your types as the first argument:
+
+```julia
+struct TheAnswer end
+
+DrWatson.wload(::Type{TheAnswer}) = TheAnswer()
+```
+"""
+function wload end
 wload(data...; kwargs...) = FileIO.load(data...; kwargs...)
+wload(::DataType, data...; kwargs...) = FileIO.load(data...; kwargs...)
 
 include("saving_files.jl")
 include("dict_list.jl")

src/saving_files.jl

@@ -36,6 +36,10 @@ end
 * `verbose = true` : print info about the process, if the file doesn't exist.
 * `wsave_kwargs = Dict()` : Keywords to pass to `wsave` (e.g. to enable
   compression).
+* `returntype::DataType = Any` : First argument passed to `wload`.
+  Set this if you defined a custom method for `wload`.
+  If `f` does not return a `returntype`, this function will issue a warning.
+  The data is still produced and saved.
 * `kwargs...` : All other keywords are propagated to `savename`.
 """
 produce_or_load(c, f; kwargs...) = produce_or_load("", c, f; kwargs...)
@@ -47,14 +51,15 @@ function produce_or_load(path, c, f::Function;
         gitpath = projectdir(), loadfile = true,
         storepatch::Bool = get(ENV, "DRWATSON_STOREPATCH", false),
         force = false, verbose = true, wsave_kwargs = Dict(),
+        returntype::DataType=Any,
         kwargs...
     )
 
     s = joinpath(path, savename(prefix, c, suffix; kwargs...))
 
     if !force && isfile(s)
         if loadfile
-            file = wload(s)
+            file = wload(returntype, s)
             return file, s
         else
             return nothing, s
@@ -66,6 +71,9 @@ function produce_or_load(path, c, f::Function;
             verbose && @info "File $s does not exist. Producing it now..."
         end
         file = f(c)
+        if !(file isa returntype)
+            @warn "Data returned from `f` is not a `$returntype`. Saving anyway, but `produce_or_load` might not be able to load the data back."
+        end
         try
             if tag
                 tagsave(s, file; safe = false, gitpath = gitpath, storepatch = storepatch, wsave_kwargs...)

test/savefiles_tests.jl

@@ -234,6 +234,53 @@ using DataFrames
     @test !isfile(spath)
 end
 
+@testset "produce_or_load with custom type" begin
+    struct Custom
+        a
+        b
+        simulation
+    end
+
+    Custom(d::Dict{String}) = Custom(d["a"], d["b"], d["simulation"])
+
+    function DrWatson._wsave(path, c::Custom)
+        wsave(joinpath(path, "a.jld2"), c.a)
+        wsave(joinpath(path, "b.jld2"), c.b)
+        wsave(joinpath(path, "simulation.jld2"), c.simulation)
+    end
+
+    function DrWatson.wload(::Type{Custom}, path)
+        a = wload(joinpath(path, "a.jld2"))
+        b = wload(joinpath(path, "b.jld2"))
+        simulation = wload(joinpath(path, "simulation.jld2"))
+        return Custom(a, b, simulation)
+    end
+
+    Base.copy(c::Custom) = c
+
+    @test !ispath(savename(simulation))
+    sim, path = produce_or_load(simulation, Custom∘f; returntype=Custom, suffix="dir")
+    @test isdir(savename(simulation))
+    @test sim.simulation.T == T
+    @test path == savename(simulation)
+    @test sim == wload(Custom, path)
+
+    # Not specifying a return type while having an overload for `_wsave` defined issues a warning:
+    msg = r"^The return type of `f` does not match"
+    rm(savename(simulation))
+    @test !isfile(savename(simulation))
+    @test_logs (:warn, msg) sim, path = produce_or_load(simulation, Custom∘f, suffix="dir")
+    @test isfile(savename(simulation))
+    @test sim == wload(Custom, path)
+
+    # Specifying the wrong return type issues a warning:
+    rm(savename(simulation))
+    @test !isfile(savename(simulation))
+    @test_logs (:warn, msg) sim, path = produce_or_load(simulation, Custom∘f; returntype=Int, suffix="dir")
+    @test isfile(savename(simulation))
+    @test sim == wload(Custom, path)
+end
+
 ################################################################################
 #                          Backup files before saving                          #
 ################################################################################