handle colors in Dict limit printing (#37568)

JuliaLang · May 11, 2022 · 9f7aa7d · 9f7aa7d
1 parent 192e388
commit 9f7aa7d
Show file tree

Hide file tree

Showing 3 changed files with 75 additions and 25 deletions.
diff --git a/base/dict.jl b/base/dict.jl
@@ -1,26 +1,5 @@
 # This file is a part of Julia. License is MIT: https://julialang.org/license
 
-function _truncate_at_width_or_chars(str, width, chars="", truncmark="…")
-    truncwidth = textwidth(truncmark)
-    (width <= 0 || width < truncwidth) && return ""
-
-    wid = truncidx = lastidx = 0
-    for (idx, c) in pairs(str)
-        lastidx = idx
-        wid += textwidth(c)
-        wid >= width - truncwidth && truncidx == 0 && (truncidx = lastidx)
-        (wid >= width || c in chars) && break
-    end
-
-    lastidx != 0 && str[lastidx] in chars && (lastidx = prevind(str, lastidx))
-    truncidx == 0 && (truncidx = lastidx)
-    if lastidx < lastindex(str)
-        return String(SubString(str, 1, truncidx) * truncmark)
-    else
-        return String(str)
-    end
-end
-
 function show(io::IO, t::AbstractDict{K,V}) where V where K
     recur_io = IOContext(io, :SHOWN_SET => t,
                              :typeinfo => eltype(t))

diff --git a/base/show.jl b/base/show.jl
@@ -48,6 +48,57 @@ show(io::IO, ::MIME"text/plain", c::ComposedFunction) = show(io, c)
 show(io::IO, ::MIME"text/plain", c::Returns) = show(io, c)
 show(io::IO, ::MIME"text/plain", s::Splat) = show(io, s)
 
+const ansi_regex = r"\x1B(?:[@-Z\\-_]|\[[0-?]*[ -/]*[@-~])"
+# An iterator similar to `pairs` but skips over "tokens" corresponding to
+# ansi sequences
+struct IgnoreAnsiIterator
+    captures::Base.RegexMatchIterator
+end
+IgnoreAnsiIterator(s::AbstractString) =
+    IgnoreAnsiIterator(eachmatch(ansi_regex, s))
+
+Base.IteratorSize(::Type{IgnoreAnsiIterator}) = Base.SizeUnknown()
+function iterate(I::IgnoreAnsiIterator, (i, m_st)=(1, iterate(I.captures)))
+    # Advance until the next non ansi sequence
+    if m_st !== nothing
+        m, j = m_st
+        if m.offset == i
+            i += sizeof(m.match)
+            return iterate(I, (i, iterate(I.captures, j)))
+        end
+    end
+    ci = iterate(I.captures.string, i)
+    ci === nothing && return nothing
+    i_prev = i
+    (c, i) = ci
+    return (i_prev => c), (i, m_st)
+end
+
+function _truncate_at_width_or_chars(ignore_ansi::Bool, str, width, chars="", truncmark="…")
+    truncwidth = textwidth(truncmark)
+    (width <= 0 || width < truncwidth) && return ""
+    wid = truncidx = lastidx = 0
+    ignore_ansi &= match(ansi_regex, str) !== nothing
+    I = ignore_ansi ? IgnoreAnsiIterator(str) : pairs(str)
+    for (_lastidx, c) in I
+        lastidx = _lastidx
+        wid += textwidth(c)
+        if wid >= (width - truncwidth) && truncidx == 0
+            truncidx = lastidx
+        end
+        (wid >= width || c in chars) && break
+    end
+    if lastidx != 0 && str[lastidx] in chars
+        lastidx = prevind(str, lastidx)
+    end
+    truncidx == 0 && (truncidx = lastidx)
+    if lastidx < lastindex(str)
+        return String(SubString(str, 1, truncidx) * truncmark)
+    else
+        return String(str)
+    end
+end
+
 function show(io::IO, ::MIME"text/plain", iter::Union{KeySet,ValueIterator})
     isempty(iter) && get(io, :compact, false) && return show(io, iter)
     summary(io, iter)
@@ -71,7 +122,7 @@ function show(io::IO, ::MIME"text/plain", iter::Union{KeySet,ValueIterator})
 
         if limit
             str = sprint(show, v, context=io, sizehint=0)
-            str = _truncate_at_width_or_chars(str, cols, "\r\n")
+            str = _truncate_at_width_or_chars(get(io, :color, false), str, cols, "\r\n")
             print(io, str)
         else
             show(io, v)
@@ -129,15 +180,15 @@ function show(io::IO, ::MIME"text/plain", t::AbstractDict{K,V}) where {K,V}
         end
 
         if limit
-            key = rpad(_truncate_at_width_or_chars(ks[i], keylen, "\r\n"), keylen)
+            key = rpad(_truncate_at_width_or_chars(get(recur_io, :color, false), ks[i], keylen, "\r\n"), keylen)
         else
             key = sprint(show, k, context=recur_io_k, sizehint=0)
         end
         print(recur_io, key)
         print(io, " => ")
 
         if limit
-            val = _truncate_at_width_or_chars(vs[i], cols - keylen, "\r\n")
+            val = _truncate_at_width_or_chars(get(recur_io, :color, false), vs[i], cols - keylen, "\r\n")
             print(io, val)
         else
             show(recur_io_v, v)
@@ -181,7 +232,7 @@ function show(io::IO, ::MIME"text/plain", t::AbstractSet{T}) where T
 
         if limit
             str = sprint(show, v, context=recur_io, sizehint=0)
-            print(io, _truncate_at_width_or_chars(str, cols, "\r\n"))
+            print(io, _truncate_at_width_or_chars(get(io, :color, false), str, cols, "\r\n"))
         else
             show(recur_io, v)
         end

diff --git a/test/dict.jl b/test/dict.jl
@@ -368,6 +368,26 @@ end
     close(io)
 end
 
+
+struct RainBowString
+    s::String
+end
+
+function Base.show(io::IO, rbs::RainBowString)
+    for s in rbs.s
+        _, color = rand(Base.text_colors)
+        print(io, color, s, "\e[0m")
+    end
+end
+
+@testset "Display with colors" begin
+    d = Dict([randstring(8) => [RainBowString(randstring(8)) for i in 1:10] for j in 1:5]...)
+    str = sprint(io -> show(io, MIME("text/plain"), d); context = (:displaysize=>(30,80), :color=>true, :limit=>true))
+    lines = split(str, '\n')
+    @test all(endswith('…'), lines[2:end])
+    @test all(x -> length(x) > 100, lines[2:end])
+end
+
 @testset "Issue #15739" begin # Compact REPL printouts of an `AbstractDict` use brackets when appropriate
     d = Dict((1=>2) => (3=>45), (3=>10) => (10=>11))
     buf = IOBuffer()