Add IO output

src/PProf.jl

@@ -45,7 +45,7 @@ using Base.StackTraces: StackFrame
 # - Mappings
 
 """
-    pprof([data, [lidict]];
+    pprof([[io], data, [lidict]];
             web = true, webhost = "localhost", webport = 57599,
             out = "profile.pb.gz", from_c = true, full_signatures = true, drop_frames = "",
             keep_frames = "", ui_relative_percentages = true, sampling_delay = nothing,
@@ -85,18 +85,59 @@ You can also use `PProf.refresh(file="...")` to open a new file in the server.
 - `ui_relative_percentages`: Passes `-relative_percentages` to pprof. Causes nodes
   ignored/hidden through the web UI to be ignored from totals when computing percentages.
 """
+function pprof end
+
+function pprof(io::IO, 
+               data::Union{Nothing, Vector{UInt}} = nothing,
+               lidict::Union{Nothing, Dict} = nothing;
+               kwargs...)
+    prof = __pprof(data, lidict; kwargs...)
+    compressed_io = GzipCompressorStream(io)
+    ProtoBuf.encode(ProtoBuf.ProtoEncoder(compressed_io), prof)
+    # can't use `close(compressed_io)` since that would close the parent `io` as well
+    write(compressed_io, CodecZlib.TranscodingStreams.TOKEN_END)
+    flush(compressed_io)
+    return nothing
+end
+
 function pprof(data::Union{Nothing, Vector{UInt}} = nothing,
                lidict::Union{Nothing, Dict} = nothing;
-               sampling_delay::Union{Nothing, UInt64} = nothing,
                web::Bool = true,
                webhost::AbstractString = "localhost",
                webport::Integer = 57599,
                out::AbstractString = "profile.pb.gz",
+               ui_relative_percentages::Bool = true,
+               kwargs...)
+    @assert !isempty(basename(out)) "`out=` must specify a file path to write to. Got unexpected: '$out'"
+    if !endswith(out, ".pb.gz")
+        out = "$out.pb.gz"
+        @info "Writing output to $out"
+    end
+
+    prof = __pprof(data, lidict; kwargs...)
+    # Write to disk
+    io = GzipCompressorStream(open(out, "w"))
+    try
+        ProtoBuf.encode(ProtoBuf.ProtoEncoder(io), prof)
+    finally
+        close(io)
+    end
+
+    if web
+        refresh(webhost = webhost, webport = webport, file = out,
+            ui_relative_percentages = ui_relative_percentages)
+    end
+
+    return out
+end
+
+function __pprof(data::Union{Nothing, Vector{UInt}} = nothing,
+               lidict::Union{Nothing, Dict} = nothing;
+               sampling_delay::Union{Nothing, UInt64} = nothing,
                from_c::Bool = true,
                full_signatures::Bool = true,
                drop_frames::Union{Nothing, AbstractString} = nothing,
                keep_frames::Union{Nothing, AbstractString} = nothing,
-               ui_relative_percentages::Bool = true,
             )
     has_meta = false
     if data === nothing
@@ -116,11 +157,6 @@ function pprof(data::Union{Nothing, Vector{UInt}} = nothing,
     if sampling_delay === nothing
         sampling_delay = ccall(:jl_profile_delay_nsec, UInt64, ())
     end
-    @assert !isempty(basename(out)) "`out=` must specify a file path to write to. Got unexpected: '$out'"
-    if !endswith(out, ".pb.gz")
-        out = "$out.pb.gz"
-        @info "Writing output to $out"
-    end
 
     string_table = OrderedDict{AbstractString, Int64}()
     enter!(string) = _enter!(string_table, string)
@@ -323,20 +359,7 @@ function pprof(data::Union{Nothing, Vector{UInt}} = nothing,
         default_sample_type = 1, # events
     )
 
-    # Write to disk
-    io = GzipCompressorStream(open(out, "w"))
-    try
-        ProtoBuf.encode(ProtoBuf.ProtoEncoder(io), prof)
-    finally
-        close(io)
-    end
-
-    if web
-        refresh(webhost = webhost, webport = webport, file = out,
-            ui_relative_percentages = ui_relative_percentages)
-    end
-
-    out
+    return prof
 end
 
 function _escape_name_for_pprof(name)