Remove hard dependency on FileIO.jl

Project.toml

@@ -6,11 +6,18 @@ version = "1.0.0"
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
-FileIO = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
 FixedPointNumbers = "53c48c17-4a7d-5ca2-90c5-79b7896eea93"
 IndirectArrays = "9b13fd28-a010-5f03-acff-a1bbcff69959"
 LeftChildRightSiblingTrees = "1d6d02ad-be62-4b6b-8a6d-2f90e265016e"
 Profile = "9abbd945-dff8-562f-b5e8-e1ebf5ef1b79"
+Requires = "ae029012-a4dd-5104-9daa-d747884805df"
+
+[weakdeps]
+FileIO = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
+Requires = "ae029012-a4dd-5104-9daa-d747884805df"
+
+[extensions]
+FlameGraphsFileIOExt = ["FileIO"]
 
 [compat]
 AbstractTrees = "0.4"
@@ -19,11 +26,13 @@ FileIO = "1.6"
 FixedPointNumbers = "0.6.1, 0.7, 0.8"
 IndirectArrays = "0.5, 1.0"
 LeftChildRightSiblingTrees = "0.2"
+Requires = "1"
 julia = "1.6"
 
 [extras]
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
+FileIO = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["InteractiveUtils", "Test"]
+test = ["InteractiveUtils", "FileIO", "Test"]

ext/FlameGraphsFileIOExt.jl

@@ -0,0 +1,38 @@
+module FlameGraphsFileIOExt
+
+isdefined(Base, :get_extension) ? (using FileIO) : (using ..FileIO)
+using Profile
+
+using FlameGraphs
+using FlameGraphs: Node, NodeData
+
+import FlameGraphs: save, load
+
+function save(f::File{format"JLPROF"}, data::AbstractVector{<:Unsigned}, lidict::Profile.LineInfoDict)
+    open(f, "w") do s
+        save(stream(s), data, lidict)
+    end
+end
+
+function save(f::File{format"JLPROF"}, g::Node{NodeData})
+    open(f, "w") do s
+        save(stream(s), g)
+    end
+end
+
+# Note: this may not work from FileIO because it returns an anonymous function for saving data
+save(f::File{format"JLPROF"}) = save(filename(f))
+
+
+function load(f::File{format"JLPROF"})
+    open(f) do s
+        skipmagic(s)
+        load(s)
+    end
+end
+
+function load(s::Stream{format"JLPROF"})
+    load(stream(s))
+end
+
+end # module FlameGraphsFileIOExt

src/FlameGraphs.jl

@@ -4,7 +4,10 @@ using Profile, LeftChildRightSiblingTrees
 using Base.StackTraces: StackFrame
 using Profile: StackFrameTree
 using Colors, FixedPointNumbers, IndirectArrays
-using FileIO
+
+@static if !isdefined(Base, :get_extension)
+    using Requires
+end
 
 # AbstractTree interface for StackFrameTree:
 using AbstractTrees
@@ -18,6 +21,12 @@ include("render.jl")
 include("sfcategory.jl")
 include("io.jl")
 
+@static if !isdefined(Base, :get_extension)
+    function __init__()
+        @require FileIO = "5789e2e9-d7fb-5bc7-8068-2c6fae9b9549" include("../ext/FlameGraphsFileIOExt.jl")
+    end
+end
+
 function _precompile_()
     ccall(:jl_generating_output, Cint, ()) == 1 || return nothing
     precompile(flametags, (Node{NodeData}, Matrix{RGB{N0f8}}))

src/io.jl

@@ -1,3 +1,5 @@
+const JLPROF_MAGIC = UInt8.(b"JLPROF\x01\x00")
+
 """
     save(f::FileIO.File)
     save(f::FileIO.File, data, lidict)
@@ -29,39 +31,47 @@ julia> @profile mapslices(sum, rand(3,3,3,3), dims=[1,2]);
 julia> save("/tmp/myprof.jlprof", Profile.retrieve()...)
 ```
 """
-function save(f::File{format"JLPROF"}, data::AbstractVector{<:Unsigned}, lidict::Profile.LineInfoDict)
+function save(filename::AbstractString, data::AbstractVector{<:Unsigned}, lidict::Profile.LineInfoDict)
+    open(filename, "w") do io
+        save(io, data, lidict)
+    end
+end
+
+function save(io::IO, data::AbstractVector{<:Unsigned}, lidict::Profile.LineInfoDict)
     data_u64 = convert(AbstractVector{UInt64}, data)
-    open(f, "w") do io
-        write(io, magic(format"JLPROF"))
-        # Write an endianness revealer
-        write(io, 0x01020304)
-        # Write an indicator that this is data/lidict format
-        write(io, 0x01)
-        write(io, Int64(length(data_u64)))
-        write(io, data_u64)
-        write(io, Int64(length(lidict)))
-        for (k, v) in lidict
-            write(io, k)
-            write(io, Int32(length(v)))
-            for sf in v
-                sfwrite(io, sf)
-            end
+    write(io, JLPROF_MAGIC)
+    # Write an endianness revealer
+    write(io, 0x01020304)
+    # Write an indicator that this is data/lidict format
+    write(io, 0x01)
+    write(io, Int64(length(data_u64)))
+    write(io, data_u64)
+    write(io, Int64(length(lidict)))
+    for (k, v) in lidict
+        write(io, k)
+        write(io, Int32(length(v)))
+        for sf in v
+            sfwrite(io, sf)
         end
     end
     return nothing
 end
 
-function save(f::File{format"JLPROF"}, g::Node{NodeData})
-    queue = Union{Nothing,typeof(g)}[]
-    open(f, "w") do io
-        write(io, magic(format"JLPROF"))
-        # Write an endianness revealer
-        write(io, 0x01020304)
-        # Write an indicator that this is node format
-        write(io, 0x02)
-        push!(queue, g)
-        savedfs!(io, queue)
+function save(filename::AbstractString, g::Node{NodeData})
+    open(filename, "w") do io
+        save(io, g)
     end
+end
+
+function save(io::IO, g::Node{NodeData})
+    queue = Union{Nothing,typeof(g)}[]
+    write(io, JLPROF_MAGIC)
+    # Write an endianness revealer
+    write(io, 0x01020304)
+    # Write an indicator that this is node format
+    write(io, 0x02)
+    push!(queue, g)
+    savedfs!(io, queue)
     return nothing
 end
 
@@ -86,8 +96,7 @@ function savedfs!(io, queue)
     return nothing
 end
 
-# Note: this doesn't work from FileIO because it returns an anonymous function for saving data
-save(f::File{format"JLPROF"}) = save(f, Profile.retrieve()...)
+save(filename::AbstractString) = save(filename, Profile.retrieve()...)
 
 """
     data, lidict = load(f::FileIO.File)
@@ -97,14 +106,17 @@ save(f::File{format"JLPROF"}) = save(f, Profile.retrieve()...)
 Load profiling data. You can reconstruct the flame graph from `flamegraph(data; lidict=lidict)`.
 Some files may already store the data in graph format, and return a single argument `g`.
 """
-function load(f::File{format"JLPROF"})
-    open(f) do io
-        skipmagic(io)
-        load(io)
-    end
+function load(filename::AbstractString)
+    open(load, filename)
 end
 
-function load(io::Stream{format"JLPROF"})
+
+function load(io::IO)
+    b0 = peek(io)
+    if b0 === JLPROF_MAGIC[1]
+        magic = read(io, length(JLPROF_MAGIC))
+        magic == JLPROF_MAGIC || error("invalid magic")
+    end
     endian = read(io, UInt32)
     endian == 0x01020304 || error("bswap not yet supported, please report as an issue to FlameGraphs.jl")
     fmt = read(io, UInt8)

test/runtests.jl

@@ -1,4 +1,4 @@
-using FlameGraphs, AbstractTrees, Colors, FileIO
+using FlameGraphs, AbstractTrees, Colors
 using Base.StackTraces: StackFrame
 using Test, Profile, InteractiveUtils
 
@@ -438,7 +438,7 @@ end
     mapslices(sum, A; dims=2)  # compile it so we don't end up profiling inference
     @profile mapslices(sum, A; dims=2)
     g = flamegraph()
-    @test FlameGraphs.depth(g) > 10
+    @test FlameGraphs.depth(g) > 7
     img = flamepixels(StackFrameCategory(), flamegraph(C=true))
     @test any(img .== colorant"orange")
     A = [1,2,3]
@@ -501,22 +501,77 @@ end
         a.from_c == b.from_c &&
         a.inlined == b.inlined
 
-    A = randn(100, 100, 200)
+    A = randn(100, 150, 200)
     Profile.clear()
     @profile mapslices(sum, A; dims=2)
-    fn = tempname()*".jlprof"
-    f = File{format"JLPROF"}(fn)
-    FlameGraphs.save(f)
-    data, lidict = FlameGraphs.load(f)
     datar, lidictr = Profile.retrieve()
-    @test data == datar
-    @test lidictr == lidict
-    rm(fn)
-    fn = tempname()*".jlprof"
-    f = File{format"JLPROF"}(fn)
-    g = flamegraph(data; lidict=lidict)
-    FlameGraphs.save(f, g)
-    gr = FlameGraphs.load(f)
-    @test nodeeq(g, gr)
-    rm(fn)
+    if datar === nothing
+        @profile mapslices(sum, A; dims=3)
+        datar, lidictr = Profile.retrieve()
+    end
+    gr = flamegraph(datar; lidict=lidictr)
+
+    @testset "w/o FileIO" begin
+        fn = tempname()*".jlprof"
+        FlameGraphs.save(fn)
+        data, lidict = FlameGraphs.load(fn)
+        @test data == datar
+        @test lidict == lidictr
+        rm(fn)
+        fn = tempname()*".jlprof"
+        FlameGraphs.save(fn, gr)
+        g = FlameGraphs.load(fn)
+        @test nodeeq(g, gr)
+        rm(fn)
+    end
+
+    @testset "w/ FileIO" begin
+        using FileIO
+        fmt = FileIO.DataFormat{:JLPROF}
+        fn = tempname()*".jlprof"
+        f = File{fmt}(fn)
+        save(f)
+        data, lidict = load(f)
+        @test data == datar
+        @test lidict == lidictr
+        rm(fn)
+        fn = tempname()*".jlprof"
+        f = File{fmt}(fn)
+        save(f, datar, lidictr)
+        data, lidict = load(f)
+        @test data == datar
+        @test lidict == lidictr
+        rm(fn)
+
+        fn = tempname()*".jlprof"
+        f = File{fmt}(fn)
+        save(f, gr)
+        g = load(f)
+        @test nodeeq(g, gr)
+        rm(fn)
+        fn = tempname()*".jlprof"
+        save(fn, gr)
+        g = load(fn)
+        @test nodeeq(g, gr)
+        rm(fn)
+
+        fn = tempname()*".non_standard_ext"
+        @test_throws Exception save(fn)
+    end
+
+    @testset "broken files" begin
+        io = IOBuffer(UInt8.(b"JLPROF\x00\x01"))
+        @test_throws Exception FlameGraphs.load(io)
+
+        fn = tempname()*".non_standard_ext"
+        write(fn, UInt8.(b"JLPROF\x00\x01"))
+        @test_throws Exception load(fn)
+        rm(fn)
+
+        io = IOBuffer(UInt8.(b"JLPROF\x01\x00\x01\x02\x03\x04"))
+        @test_throws Exception FlameGraphs.load(io)
+
+        io = IOBuffer(UInt8.(b"JLPROF\x01\x00\x04\x03\x02\x01\x03"))
+        @test_throws Exception FlameGraphs.load(io)
+    end
 end