From 50833bd1344c299b0e58123c67997e6645eeb280 Mon Sep 17 00:00:00 2001
From: TEC <git@tecosaur.net>
Date: Fri, 22 Mar 2024 01:12:51 +0800
Subject: [PATCH] Add specialised AnnotatedString pipe read/writes

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 Jamerson wasn't happy with it).

    # 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.
---
 base/strings/annotated.jl | 18 ++++++++++++++++++
 test/strings/annotated.jl |  7 +++++++
 2 files changed, 25 insertions(+)

diff --git a/base/strings/annotated.jl b/base/strings/annotated.jl
index 1ffab7a68f9e5..f14f209423a8f 100644
--- a/base/strings/annotated.jl
+++ b/base/strings/annotated.jl
@@ -500,6 +500,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})
 
diff --git a/test/strings/annotated.jl b/test/strings/annotated.jl
index f16c2bec348ca..19bf2347d8a92 100644
--- a/test/strings/annotated.jl
+++ b/test/strings/annotated.jl
@@ -209,4 +209,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