Merge pull request #91 from darsnack/darsnack/module-fix

Update to #28

Author: CarloLucibello

Project.toml

@@ -6,8 +6,9 @@ version = "0.3.1"
 julia = "0.7, 1"
 
 [extras]
+DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 Profile = "9abbd945-dff8-562f-b5e8-e1ebf5ef1b79"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "Profile"]
+test = ["DataFrames", "Test", "Profile"]

README.md

@@ -66,3 +66,24 @@ Dict{Symbol,Any} with 4 entries:
 
 This is also how the data will appear to readers in other languages, should you
 wish to move data outside of Julia.
+
+## Notes
+
+Below is some semi-official documentation on more advanced usage.
+
+### Loading custom data types within modules
+
+For packages that use BSON.jl to load data, just writing `BSON.load("mydata.bson")` will not work with custom data types. Here's a simple example of that for DataFrames.jl:
+```julia
+module A
+  using DataFrames, BSON
+  d = DataFrame(a = 1:10, b = 5:14)
+  bson("data.bson", Dict(:d=>d))
+  d2 = BSON.load("data.bson") # this will throw an error
+end
+```
+In these cases, you can specify the namespace under which to resolve types like so:
+```julia
+d2 = BSON.load("data.bson", @__MODULE__)
+```
+This will use the current module's namespace when loading the data. You could also pass any module name as the second argument (though almost all cases will use `@__MODULE__`). By default, the namespace is `Main` (i.e. the REPL).

src/anonymous.jl

@@ -50,14 +50,14 @@ end
 
 baremodule __deserialized_types__ end
 
-function newstruct_raw(cache, ::Type{TypeName}, d)
-  name = raise_recursive(d[:data][2], cache)
+function newstruct_raw(cache, ::Type{TypeName}, d, init)
+  name = raise_recursive(d[:data][2], cache, init)
   name = isdefined(__deserialized_types__, name) ? gensym() : name
   tn = ccall(:jl_new_typename_in, Ref{Core.TypeName}, (Any, Any),
              name, __deserialized_types__)
   cache[d] = tn
   names, super, parameters, types, has_instance,
-    abstr, mutabl, ninitialized = map(x -> raise_recursive(x, cache), d[:data][3:end-1])
+    abstr, mutabl, ninitialized = map(x -> raise_recursive(x, cache, init), d[:data][3:end-1])
   tn.names = names
   ndt = ccall(:jl_new_datatype, Any, (Any, Any, Any, Any, Any, Any, Cint, Cint, Cint),
               tn, tn.module, super, parameters, names, types,
@@ -68,7 +68,7 @@ function newstruct_raw(cache, ::Type{TypeName}, d)
     # use setfield! directly to avoid `fieldtype` lowering expecting to see a Singleton object already on ty
     Core.setfield!(ty, :instance, ccall(:jl_new_struct, Any, (Any, Any...), ty))
   end
-  mt = raise_recursive(d[:data][end], cache)
+  mt = raise_recursive(d[:data][end], cache, init)
   if mt != nothing
     mtname, defs, maxa, kwsorter = mt
     tn.mt = ccall(:jl_new_method_table, Any, (Any, Any), name, tn.module)

src/extensions.jl

@@ -18,9 +18,9 @@ tags[:svec] = d -> Core.svec(d[:data]...)
 
 ref(path::Symbol...) = BSONDict(:tag => "ref", :path => Base.string.([path...]))
 
-resolve(fs) = reduce((m, f) -> getfield(m, Symbol(f)), fs; init = Main)
+resolve(fs, init) = reduce((m, f) -> getfield(m, Symbol(f)), fs; init = init)
 
-tags[:ref] = d -> resolve(d[:path])
+tags[:ref] = (d, init) -> resolve(d[:path], init)
 
 function modpath(x::Module)
   y = parentmodule(x)
@@ -51,7 +51,7 @@ normalize_typeparams(x::Union{Int32,Int64}) = Int(x)
 constructtype(T, Ts) = (length(Ts) == 0) ? T : T{map(normalize_typeparams, Ts)...}
 constructtype(T::Type{Tuple}, Ts) = T{map(normalize_typeparams, Ts)...}
 
-tags[:datatype] = d -> constructtype(resolve(d[:name]), d[:params])
+tags[:datatype] = (d, init) -> constructtype(resolve(d[:name], init), d[:params])
 
 lower(v::UnionAll) =
   BSONDict(:tag => "unionall",
@@ -120,9 +120,9 @@ function newstruct(T, xs...)
   end
 end
 
-function newstruct_raw(cache, T, d)
+function newstruct_raw(cache, T, d, init)
   x = cache[d] = initstruct(T)
-  fs = map(x -> raise_recursive(x, cache), d[:data])
+  fs = map(x -> raise_recursive(x, cache, init), d[:data])
   return newstruct!(x, fs...)
 end
 
@@ -135,10 +135,10 @@ tags[:struct] = d ->
 
 iscyclic(T) = ismutable(T)
 
-raise[:struct] = function (d, cache)
-  T = d[:type] = raise_recursive(d[:type], cache)
-  iscyclic(T) || return _raise_recursive(d, cache)
-  return newstruct_raw(cache, T, d)
+raise[:struct] = function (d, cache, init)
+  T = d[:type] = raise_recursive(d[:type], cache, init)
+  iscyclic(T) || return _raise_recursive(d, cache, init)
+  return newstruct_raw(cache, T, d, init)
 end
 
 lower(v::Type{Union{}}) = BSONDict(:tag=>"jl_bottom_type")

src/read.jl

@@ -74,34 +74,38 @@ const tags = Dict{Symbol,Function}()
 
 const raise = Dict{Symbol,Function}()
 
-function _raise_recursive(d::AbstractDict, cache)
+function _raise_recursive(d::AbstractDict, cache, init)
   if haskey(d, :tag) && haskey(tags, Symbol(d[:tag]))
-    cache[d] = tags[Symbol(d[:tag])](applychildren!(x -> raise_recursive(x, cache), d))
+    if Symbol(d[:tag]) in (:ref, :datatype)
+      cache[d] = tags[Symbol(d[:tag])](applychildren!(x -> raise_recursive(x, cache, init), d), init)
+    else
+      cache[d] = tags[Symbol(d[:tag])](applychildren!(x -> raise_recursive(x, cache, init), d))
+    end
   else
     cache[d] = d
-    applychildren!(x -> raise_recursive(x, cache), d)
+    applychildren!(x -> raise_recursive(x, cache, init), d)
   end
 end
 
-function raise_recursive(d::AbstractDict, cache)
+function raise_recursive(d::AbstractDict, cache, init)
   haskey(cache, d) && return cache[d]
-  haskey(d, :tag) && haskey(raise, Symbol(d[:tag])) && return raise[Symbol(d[:tag])](d, cache)
-  _raise_recursive(d::AbstractDict, cache)
+  haskey(d, :tag) && haskey(raise, Symbol(d[:tag])) && return raise[Symbol(d[:tag])](d, cache, init)
+  _raise_recursive(d::AbstractDict, cache, init)
 end
 
-function raise_recursive(v::BSONArray, cache)
+function raise_recursive(v::BSONArray, cache, init)
   cache[v] = v
-  applychildren!(x -> raise_recursive(x, cache), v)
+  applychildren!(x -> raise_recursive(x, cache, init), v)
 end
 
-raise_recursive(x, cache) = x
+raise_recursive(x, cache, init) = x
 
-raise_recursive(x) = raise_recursive(x, IdDict())
+raise_recursive(x, init) = raise_recursive(x, IdDict(), init)
 
 parse(io::IO) = backrefs!(parse_doc(io))
 parse(path::String) = open(parse, path)
 
-load(x) = raise_recursive(parse(x))
+load(x, init=Main) = raise_recursive(parse(x), init)
 
 function roundtrip(x)
   buf = IOBuffer()

test/runtests.jl

@@ -23,6 +23,12 @@ end
 
 struct S end
 
+module A
+  using DataFrames, BSON
+  d = DataFrame(a = 1:10, b = rand(10))
+  bson("test_25_dataframe.bson", Dict(:d=>d))
+end
+
 @testset "BSON" begin
 
 @testset "Primitive Types" begin
@@ -99,4 +105,9 @@ end
   @test BSON.load(joinpath(@__DIR__, "test_MultiDimsArray_from_64bit.bson"))[:a] == ones(Float32, 2, 2)
 end
 
+@testset "Namespace other than Main #25" begin
+  @test BSON.load("test_25_dataframe.bson", A)[:d] == A.d
+  rm("test_25_dataframe.bson")
+end
+
 end