color stdlibs in Profile.print()

Author: IanButterworth

stdlib/Profile/src/Profile.jl

@@ -503,7 +503,7 @@ end
 
 # Take a file-system path and try to form a concise representation of it
 # based on the package ecosystem
-function short_path(spath::Symbol, filenamecache::Dict{Symbol, String})
+function short_path(spath::Symbol, filenamecache::Dict{Symbol, Tuple{String,String}})
     return get!(filenamecache, spath) do
         path = Base.fixup_stdlib_path(string(spath))
         if isabspath(path)
@@ -523,19 +523,19 @@ function short_path(spath::Symbol, filenamecache::Dict{Symbol, String})
                             isempty(pkgid.name) && return path # bad Project file
                             # return the joined the module name prefix and path suffix
                             path = path[nextind(path, sizeof(root)):end]
-                            return string("@", pkgid.name, path)
+                            return string("@", pkgid.name), path
                         end
                     end
                 end
             end
-            return path
+            return "", path
         elseif isfile(joinpath(Sys.BINDIR, Base.DATAROOTDIR, "julia", "base", path))
             # do the same mechanic for Base (or Core/Compiler) files as above,
             # but they start from a relative path
-            return joinpath("@Base", normpath(path))
+            return "@Base", normpath(path)
         else
             # for non-existent relative paths (such as "REPL[1]"), just consider simplifying them
-            return normpath(path) # drop leading "./"
+            return "", normpath(path) # drop leading "./"
         end
     end
 end
@@ -766,7 +766,7 @@ function flat(io::IO, data::Vector{UInt64}, lidict::Union{LineInfoDict, LineInfo
         m = m[keep]
     end
     util_perc = (1 - (nsleeping / totalshots)) * 100
-    filenamemap = Dict{Symbol,String}()
+    filenamemap = Dict{Symbol,Tuple{String,String}}()
     if isempty(lilist)
         if is_subsection
             Base.print(io, "Total snapshots: ")
@@ -790,7 +790,7 @@ end
 
 function print_flat(io::IO, lilist::Vector{StackFrame},
         n::Vector{Int}, m::Vector{Int},
-        cols::Int, filenamemap::Dict{Symbol,String},
+        cols::Int, filenamemap::Dict{Symbol,Tuple{String,String}},
         fmt::ProfileFormat)
     if fmt.sortedby === :count
         p = sortperm(n)
@@ -802,7 +802,7 @@ function print_flat(io::IO, lilist::Vector{StackFrame},
     lilist = lilist[p]
     n = n[p]
     m = m[p]
-    filenames = String[short_path(li.file, filenamemap) for li in lilist]
+    filenames = String[joinpath(short_path(li.file, filenamemap)...) for li in lilist]
     funcnames = String[string(li.func) for li in lilist]
     wcounts = max(6, ndigits(maximum(n)))
     wself = max(9, ndigits(maximum(m)))
@@ -889,7 +889,10 @@ function indent(depth::Int)
     return indent
 end
 
-function tree_format(frames::Vector{<:StackFrameTree}, level::Int, cols::Int, maxes, filenamemap::Dict{Symbol,String}, showpointer::Bool)
+# mimics Stacktraces
+const PACKAGE_FIXEDCOLORS = Dict{String, Any}("@Base" => :light_black, "@Core" => :light_black)
+
+function tree_format(frames::Vector{<:StackFrameTree}, level::Int, cols::Int, maxes, filenamemap::Dict{Symbol,Tuple{String,String}}, showpointer::Bool)
     nindent = min(cols>>1, level)
     ndigoverhead = ndigits(maxes.overhead)
     ndigcounts = ndigits(maxes.count)
@@ -924,16 +927,20 @@ function tree_format(frames::Vector{<:StackFrameTree}, level::Int, cols::Int, ma
                 else
                     fname = string(li.func)
                 end
-                filename = short_path(li.file, filenamemap)
+                pkgname, filename = short_path(li.file, filenamemap)
                 if showpointer
                     fname = string(
                         "0x",
                         string(li.pointer, base = 16, pad = 2*sizeof(Ptr{Cvoid})),
                         " ",
                         fname)
                 end
+                pkgcolor = get!(() -> popfirst!(Base.STACKTRACE_MODULECOLORS), PACKAGE_FIXEDCOLORS, pkgname)
+                remaining_path = rtruncto(filename, widthfile - length(pkgname) - 1)
                 strs[i] = string(stroverhead, "╎", base, strcount, " ",
-                    rtruncto(filename, widthfile),
+                    sprint((io,x) -> printstyled(io, x; color=pkgcolor), pkgname; context=stdout),
+                    startswith(remaining_path, "/") ? "" : "/",
+                    remaining_path,
                     ":",
                     li.line == -1 ? "?" : string(li.line),
                     "; ",
@@ -1101,7 +1108,7 @@ end
 # avoid stack overflows.
 function print_tree(io::IO, bt::StackFrameTree{T}, cols::Int, fmt::ProfileFormat, is_subsection::Bool) where T
     maxes = maxstats(bt)
-    filenamemap = Dict{Symbol,String}()
+    filenamemap = Dict{Symbol,Tuple{String,String}}()
     worklist = [(bt, 0, 0, "")]
     if !is_subsection
         Base.print(io, "Overhead ╎ [+additional indent] Count File:Line; Function\n")