From 953888eb7c8782ae1278601a0fc9d563ed7803c7 Mon Sep 17 00:00:00 2001
From: Ulrich Dobramysl <1979498+ulido@users.noreply.github.com>
Date: Fri, 3 Sep 2021 21:57:35 +0200
Subject: [PATCH] Add macro `@produce_or_load` to add source file and line tag
 (#281)

* Add macro `@produce_or_load` to add source file and line tag

Introduces a macro version of `@produce_or_load` that saves the source
file name and line number of the calling line, similar to `@tagsave`.

Adds tests that check if the appropriate file name and line number was
saved.

* Improve test coverage.

* Pass `gitpath` to `scripttag!` in `@produce_or_load`.

* Add `@produce_or_load` to docs.

* Refactor `tag!` to avoid code duplication.

* Fix test failure by making `scripttag!` argument types less restrictive.

* Increment minor version and add changelog entry.

* Replace method dispatch type checks with an assertion.
---
 CHANGELOG.md            |  2 +
 Project.toml            |  2 +-
 docs/src/save.md        |  1 +
 src/saving_files.jl     | 36 +++++++++++++++-
 src/saving_tools.jl     | 92 ++++++++++++++++++++++++++++++-----------
 test/savefiles_tests.jl | 55 ++++++++++++++++++++++++
 test/stools_tests.jl    |  5 +++
 7 files changed, 166 insertions(+), 27 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2d096039..eab7c2f9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,3 +1,5 @@
+# 2.4.0
+* Add the macro version of `produce_or_load` to enable tagging with the calling source file and line.
 # 2.3.0
 * Enable pass through of kwargs to `wsave` in `produce_or_load`, `tagsave` and `safesave` (to e.g. allow compression in JLD2 files).
 # 2.2.0
diff --git a/Project.toml b/Project.toml
index e1a7fc2e..0b2125c0 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "DrWatson"
 uuid = "634d3b9d-ee7a-5ddf-bec9-22491ea816e1"
 repo = "https://github.com/JuliaDynamics/DrWatson.jl.git"
-version = "2.3.0"
+version = "2.4.0"
 
 [deps]
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
diff --git a/docs/src/save.md b/docs/src/save.md
index 0db96564..fb30a9c4 100644
--- a/docs/src/save.md
+++ b/docs/src/save.md
@@ -56,6 +56,7 @@ In addition, it attempts to minimize computing energy spent on getting a result.
 
 ```@docs
 produce_or_load
+@produce_or_load
 istaggable
 ```
 
diff --git a/src/saving_files.jl b/src/saving_files.jl
index dfed8be9..87c15bdf 100644
--- a/src/saving_files.jl
+++ b/src/saving_files.jl
@@ -1,4 +1,4 @@
-export produce_or_load, tagsave, @tagsave, safesave
+export produce_or_load, @produce_or_load, tagsave, @tagsave, safesave
 
 """
     produce_or_load([path="",] config, f; kwargs...) -> file, s
@@ -80,6 +80,40 @@ function produce_or_load(path, c, f::Function;
     end
 end
 
+"""
+    @produce_or_load([path="",] config, f; kwargs...)
+Same as [`produce_or_load`](@ref) but one more field `:script` is added that records
+the local path of the script and line number that called `@produce_or_load`, see [`@tag!`](@ref).
+"""
+macro produce_or_load(f, path, config, args...)
+    args = Any[args...]
+    # Keywords added after a ; are moved to the front of the expression
+    # that is passed to the macro. So instead of getting the function in f
+    # an Expr is passed.
+    if f isa Expr && f.head == :parameters
+        length(args) > 0 || return :(throw(MethodError(@produce_or_load,$(esc(f)),$(esc(path)),$(esc(config)),$(esc.(args)...))))
+        extra_kw_def = f.args
+        f = path
+        path = config
+        config = popfirst!(args)
+        append!(args,extra_kw_def)
+    end
+    # Save the source file name and line number of the calling line.
+    s = QuoteNode(__source__)
+    # Wrap the function f, such that the source can be saved in the data Dict.
+    return quote
+        produce_or_load($(esc(path)), $(esc(config)), $(esc.(convert_to_kw.(args))...)) do k
+            data = $(esc(f))(k)
+            # Extract the `gitpath` kw arg if it's there
+            kws = ((;kwargs...)->Dict(kwargs...))($(esc.(convert_to_kw.(args))...))
+            gitpath = get(kws, :gitpath, projectdir())
+            # Include the script tag with checking for the type of dict keys, etc.
+            scripttag!(data, $s; gitpath = gitpath)
+            data
+        end
+    end
+end
+
 ################################################################################
 #                             tag saving                                       #
 ################################################################################
diff --git a/src/saving_tools.jl b/src/saving_tools.jl
index 9405e253..0fcd2180 100644
--- a/src/saving_tools.jl
+++ b/src/saving_tools.jl
@@ -193,44 +193,86 @@ Dict{Symbol,Any} with 3 entries:
 ```
 """
 function tag!(d::Dict{K,T}; gitpath = projectdir(), storepatch = true, force = false, source = nothing) where {K,T}
+    @assert (K <: Union{Symbol,String}) "We only know how to tag dictionaries that have keys that are strings or symbols"
     c = gitdescribe(gitpath)
-    patch = gitpatch(gitpath)
-    @assert (Symbol <: K) || (String <: K)
-    if K == Symbol
-        commitname, patchname, scriptname = :gitcommit, :gitpatch, :script
-    else
-        commitname, patchname, scriptname = "gitcommit", "gitpatch", "script"
-    end
-
     c === nothing && return d # gitpath is not a git repo
+
+    # Get the appropriate keys
+    commitname = keyname(d, :gitcommit)
+    patchname = keyname(d, :gitpatch)
+
     if haskey(d, commitname) && !force
         @warn "The dictionary already has a key named `gitcommit`. We won't "*
         "add any Git information."
-        return d
-    end
-    if String <: T
-        d[commitname] = c
-        if storepatch && (patch != nothing)
-            d[patchname] = patch
-        end
     else
-        d = Dict{K, promote_type(T, String)}(d)
+        d = checktagtype!(d) 
         d[commitname] = c
-        if patch!=""
-            d[patchname] = patch
+        # Only include patch info if `storepatch` is true and if we can get the info.
+        if storepatch
+            patch = gitpatch(gitpath)
+            if (patch != nothing) && (patch != "")
+                d[patchname] = patch
+            end
         end
     end
-    if source !== nothing && !force
-        if haskey(d, scriptname)
-            @warn "The dictionary already has a key named `script`. We won't "*
-            "overwrite it with the script name."
-        else
-            d[scriptname] = relpath(sourcename(source), gitpath)
-        end
+
+    # Include source file and line number info if given.
+    if source !== nothing
+        d = scripttag!(d, source; gitpath = gitpath, force = force)
     end
+
     return d
 end
 
+"""
+    keyname(d::Dict{K,T}, key) where {K<:Union{Symbol,String},T}
+
+Check the key type of `d` and convert `key` to the appropriate type.
+"""
+function keyname(d::Dict{K,T}, key) where {K<:Union{Symbol,String},T}
+    if K == Symbol
+        return Symbol(key)
+    end
+    return String(key)
+end
+
+"""
+    checktagtype!(d::Dict{K,T}) where {K<:Union{Symbol,String},T}
+
+Check if the value type of `d` allows `String` and promote it to do so if not.
+"""
+function checktagtype!(d::Dict{K,T}) where {K<:Union{Symbol,String},T}
+    if !(String <: T)
+        d = Dict{K, promote_type(T, String)}(d)
+    end
+    d
+end
+
+"""
+    scripttag!(d::Dict{K,T}, source::LineNumberNode; gitpath = projectdir(), force = false) where {K<:Union{Symbol,String},T}
+
+Include a `script` field in `d`, containing the source file and line number in
+`source`. Do nothing if the field is already present unless `force = true`. Uses
+`gitpath` to make the source file path relative.
+"""
+function scripttag!(d::Dict{K,T}, source; gitpath = projectdir(), force = false) where {K,T}
+    # We want this functionality to be separate from `tag!` to allow
+    # inclusion of this information without the git tagging
+    # functionality.
+    # To be used in `tag!` and `@produce_or_load`.
+    # We have to assert the key type here again because `scripttag!` can be called
+    # from `@produce_or_load` without going through `tag!`.
+    @assert (K <: Union{Symbol,String}) "We only know how to tag dictionaries that have keys that are strings or symbols"
+    scriptname = keyname(d, :script)
+    if haskey(d, scriptname) && !force
+        @warn "The dictionary already has a key named `script`. We won't "*
+            "overwrite it with the script name."
+    else
+        d = checktagtype!(d)
+        d[scriptname] = relpath(sourcename(source), gitpath)
+    end
+    return d
+end
 sourcename(s) = string(s)
 sourcename(s::LineNumberNode) = string(s.file)*"#"*string(s.line)
 
diff --git a/test/savefiles_tests.jl b/test/savefiles_tests.jl
index 41e44888..70df8489 100644
--- a/test/savefiles_tests.jl
+++ b/test/savefiles_tests.jl
@@ -105,6 +105,37 @@ end
     rm(savename(simulation, ending))
     @test !isfile(savename(simulation, ending))
 
+    @test !isfile(savename(simulation, ending))
+    # Produce and save data, preserve source file name and line for test below.
+    # Line needs to be saved on the same line as produce_or_load!
+    sim, path = @produce_or_load(f, "", simulation, suffix = ending); fname = @__FILE__; line = @__LINE__
+    @test isfile(savename(simulation, ending))
+    @test sim["simulation"].T == T
+    @test path == savename(simulation, ending)
+    sim, path = @produce_or_load(f, "", simulation, suffix = ending)
+    @test sim["simulation"].T == T
+    # Test if source was included and that the file name and line number matches the first invocation
+    # (and not the second!)
+    @test "script" ∈ keys(sim)
+    @test sim["script"] |> typeof == String
+    @test sim["script"] == joinpath(relpath(fname, projectdir()) * "#$(line)")
+    rm(savename(simulation, ending))
+    @test !isfile(savename(simulation, ending))
+
+    # Test if tag = true does not interfere with macro script tagging.
+    # Use a semicolon before the `suffix` keyword to test that code path as well.
+    sim, path = @produce_or_load(f, "", simulation, tag = true; suffix = ending); fname = @__FILE__; line = @__LINE__
+    sim, path = @produce_or_load(f, "", simulation; suffix = ending)
+    # Test if source was included and that the file name and line number matches the first invocation
+    # (and not the second!)
+    @test sim["script"] == joinpath(relpath(fname, projectdir()) * "#$(line)")
+    rm(savename(simulation, ending))
+
+    # Test that the internal function `scripttag!` properly warns if the Dict already has a `script` key.
+    # This also tests the case where the `Dict` has a `Symbol` key type.
+    @test_logs((:warn, "The dictionary already has a key named `script`. We won't overwrite it with the script name."),
+               DrWatson.scripttag!(Dict(:script => "test"), LineNumberNode(1)))
+
     @test !isfile(savename(simulation, ending))
     sim, path = produce_or_load("", simulation; suffix = ending) do simulation
         @test typeof(simulation.T) <: Real
@@ -158,6 +189,30 @@ end
     # Leave no trace
     rm(sn_uncomp)
     rm(sn_comp)
+
+    # Check the macro version
+    sn_uncomp = savename(Dict("compress" => false), "jld2")
+    sn_comp = savename(Dict("compress" => true), "jld2")
+    # Files cannot exist yet
+    @test !isfile(sn_uncomp)
+    @test !isfile(sn_comp)
+    for compress in [false, true]
+        wsave_kwargs = Dict(:compress => compress)
+        @produce_or_load("", wsave_kwargs, suffix = "jld2", wsave_kwargs=wsave_kwargs) do c
+            data
+        end
+    end
+    # Check if both files exist now
+    @test isfile(sn_uncomp)
+    @test isfile(sn_comp)
+    # Test if the compressed file is smaller
+    size_uncomp = filesize(sn_uncomp)
+    size_comp = filesize(sn_comp)
+    @test size_uncomp > size_comp
+    # Leave no trace
+    rm(sn_uncomp)
+    rm(sn_comp)
+
 end
 
 @test produce_or_load(simulation, f; loadfile = false)[1] == nothing
diff --git a/test/stools_tests.jl b/test/stools_tests.jl
index a0298573..a3895c5e 100644
--- a/test/stools_tests.jl
+++ b/test/stools_tests.jl
@@ -27,6 +27,11 @@ for d in (d1, d2)
     @test d[keytype(d)(:gitcommit)] |> typeof <: String
 end
 
+# Test assertion error when the data has a incompatible key type
+@test_throws AssertionError("We only know how to tag dictionaries that have keys that are strings or symbols") tag!(Dict{Int64,Any}(1 => 2))
+@test_throws AssertionError("We only know how to tag dictionaries that have keys that are strings or symbols") DrWatson.scripttag!(Dict{Int64,Any}(1 => 2), "foo")
+
+
 # @tag!
 for d in (d1, d2)
     d = @tag!(d, gitpath=@__DIR__)