Add specialised AnnotatedString pipe read/writes (#53715)

Ensure that when an AnnotatedIOBuffer is wrapped in an IOContext (or
similar AnnotatedPipe-based construct), that writes of annotated
strings/chars and reading out an AnnotatedString is unimpeded by the
IOContext wrapping.

Without these specialisations, the generic pipe_reader/pipe_writer
fallbacks will directly access the underlying IOBuffer and annotations
will be lost.

There are a number of scenarios in which one might want to combine an
AnnotatedIOBuffer and IOContext (for example setting the compact
property). Losing annotations in such scenarios is highly undesirable.
It is of particular note that this can arise in situations where you
can't unwrap the IOContext as needed, for example when passing IO to a
function you do not control (which is currently extremely hard to work
around).

Getting this right is a little difficult, and a few approaches have been
tried. Initially, I added IOContext{AnnotatedIOBuffer} specialisations
to show.jl, but arguably it's a bit of a code smell to specialise in
this way (and Jameson wasn't happy with it, with concerns that it could
be tricked by IOContext{Any}).

    # So that read/writes with `IOContext` (and any similar `AbstractPipe` wrappers)
    # work as expected.
    write(io::IOContext{AnnotatedIOBuffer}, s::Union{AnnotatedString, SubString{<:AnnotatedString}}) =
        write(io.io, s)
    write(io::AnnotatedIOBuffer, c::AnnotatedChar) = write(io.io, c)

Then I tried making it so that IOContext writes dispatched on the
wrapped IO type, but of course that broke cases like IOContext{IOBuffer}
with :color=>true.

    # So that read/writes with `IOContext` (and any similar `AbstractPipe` wrappers)
    # work as expected.
    write(io::AbstractPipe, s::Union{AnnotatedString, SubString{<:AnnotatedString}}) =
        write(pipe_writer(io), s)
    write(io::AbstractPipe, c::AnnotatedChar) = write(pipe_writer(io), c)

Finally, we have the current AbstractPipe + Annotated type
specialisation, which IOContext is just an instance of. To avoid
behaving too broadly, we need to confirm that the underlying IO is
actually an AnnotatedIOBuffer. I'm still not happy with this, only idea
I've had other than implementing IOContext{AnnotatedIOBuffer} methods
that actually seems viable, and I've had trouble soliciting help from
other people brainstorming here.

If somebody can implement something cleaner here in the future, I'd be
thrilled.

Author: tecosaur

base/strings/annotated.jl

@@ -518,6 +518,24 @@ function write(dest::AnnotatedIOBuffer, src::AnnotatedIOBuffer)
     nb
 end
 
+# So that read/writes with `IOContext` (and any similar `AbstractPipe` wrappers)
+# work as expected.
+function write(io::AbstractPipe, s::Union{AnnotatedString, SubString{<:AnnotatedString}})
+    if pipe_writer(io) isa AnnotatedIOBuffer
+        write(pipe_writer(io), s)
+    else
+        invoke(write, Tuple{IO, typeof(s)}, io, s)
+    end::Int
+end
+# Can't be part of the `Union` above because it introduces method ambiguities
+function write(io::AbstractPipe, c::AnnotatedChar)
+    if pipe_writer(io) isa AnnotatedIOBuffer
+        write(pipe_writer(io), c)
+    else
+        invoke(write, Tuple{IO, typeof(c)}, io, c)
+    end::Int
+end
+
 """
     _clear_annotations_in_region!(annotations::Vector{Tuple{UnitRange{Int}, Pair{Symbol, Any}}}, span::UnitRange{Int})
 

test/strings/annotated.jl

@@ -213,4 +213,11 @@ end
         @test write(aio2, Base.AnnotatedChar('c', [:b => 2, :c => 3, :d => 4])) == 1
         @test Base.annotations(aio2) == [(1:2, :a => 1), (1:3, :b => 2), (3:3, :c => 3), (3:3, :d => 4)]
     end
+    # Working through an IOContext
+    aio = Base.AnnotatedIOBuffer()
+    wrapio = IOContext(aio)
+    @test write(wrapio, Base.AnnotatedString("hey", [(1:3, :x => 1)])) == 3
+    @test write(wrapio, Base.AnnotatedChar('a', [:y => 2])) == 1
+    @test read(seekstart(aio), Base.AnnotatedString) ==
+        Base.AnnotatedString("heya", [(1:3, :x => 1), (4:4, :y => 2)])
 end