Support BSON and JLD2 file formats

Project.toml

@@ -6,11 +6,13 @@ version = "0.1.5"
 [deps]
 BSON = "fbb218c0-5317-5bc6-957e-2ee96dd4b1f0"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
+JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 
 [compat]
 BSON = "0.3.4"
 Dates = "<0.0.1, 1.0"
+JLD2 = "0.4, 0.5, 0.6"
 Logging = "<0.0.1, 1.0"
 julia = "1.0"
 

README.md

@@ -24,6 +24,9 @@ Subsequent runs load the saved output from the file `test.bson`
 rather than re-running the potentially time-consuming computations!
 Especially handy for long simulations.
 
+The file format is determined by the file extension:
+`.bson` for [BSON.jl](https://github.com/JuliaIO/BSON.jl) and `.jld2` for [JLD2.jl](https://github.com/JuliaIO/JLD2.jl).
+
 An example of the output:
 
 ```julia
@@ -78,6 +81,9 @@ Subsequent runs load the saved values from the file `test.bson`
 rather than re-running the potentially time-consuming computations!
 Especially handy for long simulations.
 
+As with the function form, the file format is determined by the file extension:
+`.bson` for [BSON.jl](https://github.com/JuliaIO/BSON.jl) and `.jld2` for [JLD2.jl](https://github.com/JuliaIO/JLD2.jl).
+
 An example of the output:
 
 ```julia
@@ -314,6 +320,22 @@ julia> cache(SUBSET === Colon() ? "fullsweep.bson" : nothing) do
 ```
 Note that the full cache was not generated here.
 
+## File formats
+
+CacheVariables.jl supports two file formats:
+
+- [BSON.jl](https://github.com/JuliaIO/BSON.jl): use a filename with extension `.bson`
+- [JLD2.jl](https://github.com/JuliaIO/JLD2.jl): use a filename with extension `.jld2`
+
+Optional format-specific arguments can be passed as keyword arguments.
+In particular, for BSON files, you can pass the `mod` keyword to specify the module context for loading:
+```julia
+cache("data.bson"; mod = @__MODULE__) do
+    # your computation
+end
+```
+This can be useful when working in modules or in Pluto notebooks (see, e.g., [https://github.com/JuliaIO/BSON.jl?tab=readme-ov-file#loading-custom-data-types-within-modules](https://github.com/JuliaIO/BSON.jl?tab=readme-ov-file#loading-custom-data-types-within-modules)).
+
 ## Related packages
 
 - [Memoization.jl](https://github.com/marius311/Memoization.jl)

src/CacheVariables.jl

@@ -1,6 +1,7 @@
 module CacheVariables
 
 using BSON
+using JLD2
 import Dates
 import Logging: @info
 
@@ -17,10 +18,12 @@ end
 """
     @cache path code [overwrite]
 
-Cache results from running `code` using BSON file at `path`.
+Cache results from running `code` using the file at `path`.
 Load if the file exists; run and save if it does not.
 Run and save either way if `overwrite` is true (default is false).
 
+The file format is determined by the file extension: `.bson` or `.jld2`.
+
 Tip: Use `begin...end` for `code` to cache blocks of code.
 
 Caveat: The variable name `ans` is used for storing the final output,
@@ -72,23 +75,63 @@ macro cache(path, ex::Expr, overwrite = false)
             _msg = ispath($(esc(path))) ? "Overwriting " : "Saving to "
             @info(string(_msg, $(esc(path)), "\n", $(vardesc)))
             mkpath(splitdir($(esc(path)))[1])
-            bson($(esc(path)); $(esc(varlist))..., ans = _ans)
+            # Determine format by extension
+            if endswith($(esc(path)), ".bson")
+                # Use BSON.bson to maintain backward compatibility (Symbol keys)
+                bson($(esc(path)); $(esc(varlist))..., ans = _ans)
+            elseif endswith($(esc(path)), ".jld2")
+                # Use JLD2 for better type support
+                # Note: JLD2 converts Symbol keys to String keys when saving
+                _data = Dict{Symbol,Any}()
+                $([:(_data[$(QuoteNode(var))] = $(esc(var))) for var in vars]...)
+                _data[:ans] = _ans
+                JLD2.jldsave($(esc(path)); _data...)
+            else
+                error("Unsupported file extension for $($(esc(path))). Only .bson and .jld2 are supported.")
+            end
             _ans
         else
             @info(string("Loading from ", $(esc(path)), "\n", $(vardesc)))
-            data = BSON.load($(esc(path)), @__MODULE__)
-            $(esc(vartuple)) = getindex.(Ref(data), $vars)
-            data[:ans]
+            # Determine format by extension
+            if endswith($(esc(path)), ".bson")
+                _data = BSON.load($(esc(path)), @__MODULE__)
+                # BSON uses Symbol keys
+                _vars_syms = $vars
+                $(esc(vartuple)) = getindex.(Ref(_data), _vars_syms)
+                _data[:ans]
+            elseif endswith($(esc(path)), ".jld2")
+                _data = JLD2.load($(esc(path)))
+                # JLD2 saves with String keys (even when we pass Symbol keys)
+                _vars_strs = $(map(String, vars))
+                $(esc(vartuple)) = map(_vars_strs) do str
+                    if haskey(_data, str)
+                        _data[str]
+                    else
+                        error("Cache file $($(esc(path))) missing required key: $str")
+                    end
+                end
+                if haskey(_data, "ans")
+                    _data["ans"]
+                else
+                    error("Cache file $($(esc(path))) missing required key: ans")
+                end
+            else
+                error("Unsupported file extension for $($(esc(path))). Only .bson and .jld2 are supported.")
+            end
         end
     end
 end
 
 """
-    cache(f, path; mod = @__MODULE__)
+    cache(f, path; kwargs...)
 
-Cache output from running `f()` using BSON file at `path`.
+Cache output from running `f()` using the file at `path`.
 Load if the file exists; run and save if it does not.
-Use `mod` keyword argument to specify module.
+
+The file format is determined by the file extension: `.bson` or `.jld2`.
+
+Additional keyword arguments are passed to the underlying save/load functions.
+For BSON files, you can pass `mod = @__MODULE__` to specify the module for loading.
 
 Tip: Use `do...end` to cache output from a block of code.
 
@@ -109,32 +152,69 @@ julia> cache("test.bson") do
        end
 [ Info: Loading from test.bson
 (a = "a very time-consuming quantity to compute", b = "a very long simulation to run")
+
+julia> cache("test.jld2") do
+         return big"123456789012345678901234567890"
+       end
+[ Info: Saving to test.jld2
+123456789012345678901234567890
 ```
 """
-function cache(@nospecialize(f), path; mod = @__MODULE__)
+function cache(@nospecialize(f), path; kwargs...)
     if !ispath(path)
         ans = f()
         @info(string("Saving to ", path, "\n"))
         mkpath(splitdir(path)[1])
-        bson(path; ans = ans)
+        # Determine format by extension
+        if endswith(path, ".bson")
+            # Use BSON.jl directly to maintain compatibility
+            bson(path; ans = ans)
+        elseif endswith(path, ".jld2")
+            # Use JLD2 for better type support
+            JLD2.jldsave(path; ans = ans)
+        else
+            error("Unsupported file extension for $path. Only .bson and .jld2 are supported.")
+        end
         return ans
     else
         @info(string("Loading from ", path, "\n"))
-        data = BSON.load(path, mod)
-        return data[:ans]
+        # Determine format by extension
+        if endswith(path, ".bson")
+            # For BSON files with mod argument, use BSON.load directly
+            if haskey(kwargs, :mod)
+                data = BSON.load(path, kwargs[:mod])
+            else
+                data = BSON.load(path)
+            end
+            return data[:ans]
+        elseif endswith(path, ".jld2")
+            # Use JLD2
+            data = JLD2.load(path)
+            if haskey(data, "ans")
+                return data["ans"]
+            else
+                error("Cache file $path does not contain required 'ans' key. File may be corrupted or incompatible.")
+            end
+        else
+            error("Unsupported file extension for $path. Only .bson and .jld2 are supported.")
+        end
     end
 end
-function cache(@nospecialize(f), ::Nothing; mod = @__MODULE__)
+function cache(@nospecialize(f), ::Nothing; kwargs...)
     @info("No path provided, running without caching.")
     return f()
 end
 
 """
-    cachemeta(f, path; mod = @__MODULE__)
+    cachemeta(f, path; kwargs...)
 
-Cache output from running `f()` using BSON file at `path` with additional metadata.
+Cache output from running `f()` using the file at `path` with additional metadata.
 Load if the file exists; run and save if it does not.
-Use `mod` keyword argument to specify module.
+
+The file format is determined by the file extension: `.bson` or `.jld2`.
+
+Additional keyword arguments are passed to the underlying save/load functions.
+For BSON files, you can pass `mod = @__MODULE__` to specify the module for loading.
 
 Saves and displays the following metadata:
 - Julia version (from `VERSION`)
@@ -164,8 +244,8 @@ julia> cachemeta("test.bson") do
 (a = "a very time-consuming quantity to compute", b = "a very long simulation to run")
 ```
 """
-function cachemeta(@nospecialize(f), path; mod = @__MODULE__)
-    version, whenrun, runtime, ans = cache(path; mod = mod) do
+function cachemeta(@nospecialize(f), path; kwargs...)
+    version, whenrun, runtime, ans = cache(path; kwargs...) do
         version = VERSION
         whenrun = Dates.now(Dates.UTC)
         runtime = @elapsed ans = f()