From 3c2f0e2fb67fb5b7d4dcec758e816422fe07f7ee Mon Sep 17 00:00:00 2001
From: TEC <git@tecosaur.net>
Date: Sat, 21 Oct 2023 19:48:18 +0800
Subject: [PATCH] Introduce AnnotatedIOBuffer

This allows for styled content to be constructed incrementally, without
resorting to repeated concatenation. It operates very similarly to
IOContext, just with a special `write` method and specifically wrapping
an IOBuffer.
---
 base/show.jl              |  3 ++
 base/strings/annotated.jl | 59 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 62 insertions(+)

diff --git a/base/show.jl b/base/show.jl
index 079bf5a423cc0d..f456426046efff 100644
--- a/base/show.jl
+++ b/base/show.jl
@@ -429,6 +429,9 @@ function show_circular(io::IOContext, @nospecialize(x))
     return false
 end
 
+AnnotatedIOBuffer(io::IOContext{IOBuffer}) = AnnotatedIOBuffer(io.io, io.dict)
+IOContext(io::AnnotatedIOBuffer) = IOContext(io.io, io.dict)
+
 """
     show([io::IO = stdout], x)
 
diff --git a/base/strings/annotated.jl b/base/strings/annotated.jl
index c0f3623f416963..95270524221f3b 100644
--- a/base/strings/annotated.jl
+++ b/base/strings/annotated.jl
@@ -386,3 +386,62 @@ annotations(s::SubString{<:AnnotatedString}, pos::UnitRange{<:Integer}) =
 Get all annotations of `chr`.
 """
 annotations(c::AnnotatedChar) = c.annotations
+
+## AnnotatedIOBuffer
+
+struct AnnotatedIOBuffer <: IO
+    io::IOBuffer
+    annotations::Vector{Tuple{UnitRange{Int}, Pair{Symbol, Any}}}
+    dict::ImmutableDict{Symbol, Any}
+end
+
+AnnotatedIOBuffer(io::IOBuffer, dict::ImmutableDict = ImmutableDict{Symbol,Any}()) =
+    AnnotatedIOBuffer(io, Vector{Tuple{UnitRange{Int}, Pair{Symbol, Any}}}(), dict)
+
+function AnnotatedIOBuffer(io::IOBuffer, (k, v)::Pair{Symbol, Any}, kvs::Pair{Symbol, Any}...)
+    dict = ImmutableDict{Symbol,Any}(k, v)
+    for (k, v) in kvs
+        dict = ImmutableDict{Symbol,Any}(dict, k, v)
+    end
+    AnnotatedIOBuffer(io, dict)
+end
+
+AnnotatedIOBuffer() = AnnotatedIOBuffer(IOBuffer())
+
+function show(io::IO, annio::AnnotatedIOBuffer)
+    show(io, AnnotatedIOBuffer)
+    print(io, '(', annio.io.size, " bytes)")
+end
+
+position(io::AnnotatedIOBuffer) = position(io.io)
+lock(io::AnnotatedIOBuffer) = lock(io.io)
+unlock(io::AnnotatedIOBuffer) = unlock(io.io)
+get(io::AnnotatedIOBuffer, key, default) = key == :color || get(io.dict, key, default)
+displaysize(io::AnnotatedIOBuffer) =
+    if haskey(io.dict, :displaysize) io.dict[:displaysize]::Tuple{Int,Int} else displaysize(io.io) end
+
+function write(io::AnnotatedIOBuffer, astr::Union{AnnotatedString, SubString{<:AnnotatedString}})
+    astr = AnnotatedString(astr)
+    offset = position(io.io)
+    for (region, annot) in astr.annotations
+        start, stop = first(region), last(region)
+        push!(io.annotations, (start+offset:stop+offset, annot))
+    end
+    write(io.io, astr)
+end
+write(io::AnnotatedIOBuffer, achr::AnnotatedChar) = write(aio, AnnotatedString(achr))
+write(io::AnnotatedIOBuffer, x::AbstractString) = write(io.io, x)
+write(io::AnnotatedIOBuffer, x::UInt8) = write(io.io, x)
+write(io::Base.AnnotatedIOBuffer, x::AbstractString) = write(io.io, x)
+write(io::IO, s::Union{SubString{String}, String}) = write(io.io, x)
+
+function take!(aiob::AnnotatedIOBuffer)
+    str = String(take!(aiob.io))
+    annot = copy(aiob.annotations)
+    empty!(aiob.annotations)
+    seekstart(aiob.io)
+    str, annot
+end
+
+AnnotatedString((str, annots)::Tuple{<:AbstractString, Vector{Tuple{UnitRange{Int}, Pair{Symbol, Any}}}}) =
+    AnnotatedString(str, annots)