Add annotate! method for AnnotatedIOBuffer (#53284)

The annotate! function provides a convenient way of adding annotations
to an AnnotatedString/AnnotatedChar without accessing any of the
implementation details of the Annotated* types.

When AnnotatedIOBuffer was introduced, an appropriate annotations method
was added, but annotate! was missed. To correct that, we refactor the
current annotate! method for AnnotatedString to a more general helper
function _annotate! that operates on the annotation vector itself, and
use this new helper method to easily provide an annotate! method for
AnnotatedIOBuffer.

Author: tecosaur

base/strings/annotated.jl

@@ -316,28 +316,30 @@ reverse(s::SubString{<:AnnotatedString}) = reverse(AnnotatedString(s))
 
 ## End AbstractString interface ##
 
-"""
-    annotate!(str::AnnotatedString, [range::UnitRange{Int}], label::Symbol => value)
-    annotate!(str::SubString{AnnotatedString}, [range::UnitRange{Int}], label::Symbol => value)
-
-Annotate a `range` of `str` (or the entire string) with a labeled value (`label` => `value`).
-To remove existing `label` annotations, use a value of `nothing`.
-"""
-function annotate!(s::AnnotatedString, range::UnitRange{Int}, @nospecialize(labelval::Pair{Symbol, <:Any}))
+function _annotate!(annlist::Vector{Tuple{UnitRange{Int}, Pair{Symbol, Any}}}, range::UnitRange{Int}, @nospecialize(labelval::Pair{Symbol, <:Any}))
     label, val = labelval
     if val === nothing
-        indices = searchsorted(s.annotations, (range,), by=first)
-        labelindex = filter(i -> first(s.annotations[i][2]) === label, indices)
+        indices = searchsorted(annlist, (range,), by=first)
+        labelindex = filter(i -> first(annlist[i][2]) === label, indices)
         for index in Iterators.reverse(labelindex)
-            deleteat!(s.annotations, index)
+            deleteat!(annlist, index)
         end
     else
-        sortedindex = searchsortedlast(s.annotations, (range,), by=first) + 1
-        insert!(s.annotations, sortedindex, (range, Pair{Symbol, Any}(label, val)))
+        sortedindex = searchsortedlast(annlist, (range,), by=first) + 1
+        insert!(annlist, sortedindex, (range, Pair{Symbol, Any}(label, val)))
     end
-    s
 end
 
+"""
+    annotate!(str::AnnotatedString, [range::UnitRange{Int}], label::Symbol => value)
+    annotate!(str::SubString{AnnotatedString}, [range::UnitRange{Int}], label::Symbol => value)
+
+Annotate a `range` of `str` (or the entire string) with a labeled value (`label` => `value`).
+To remove existing `label` annotations, use a value of `nothing`.
+"""
+annotate!(s::AnnotatedString, range::UnitRange{Int}, @nospecialize(labelval::Pair{Symbol, <:Any})) =
+    (_annotate!(s.annotations, range, labelval); s)
+
 annotate!(ss::AnnotatedString, @nospecialize(labelval::Pair{Symbol, <:Any})) =
     annotate!(ss, firstindex(ss):lastindex(ss), labelval)
 
@@ -426,6 +428,9 @@ copy(io::AnnotatedIOBuffer) = AnnotatedIOBuffer(copy(io.io), copy(io.annotations
 
 annotations(io::AnnotatedIOBuffer) = io.annotations
 
+annotate!(io::AnnotatedIOBuffer, range::UnitRange{Int}, @nospecialize(labelval::Pair{Symbol, <:Any})) =
+    (_annotate!(io.annotations, range, labelval); io)
+
 function write(io::AnnotatedIOBuffer, astr::Union{AnnotatedString, SubString{<:AnnotatedString}})
     astr = AnnotatedString(astr)
     offset = position(io.io)

test/strings/annotated.jl

@@ -115,6 +115,12 @@ end
     @test write(aio, ' ') == 1
     @test write(aio, Base.AnnotatedString("world", [(1:5, :tag => 2)])) == 5
     @test Base.annotations(aio) == [(1:5, :tag => 1), (7:11, :tag => 2)]
+    # Check `annotate!`, including region sorting
+    @test truncate(aio, 0).io.size == 0
+    @test write(aio, "hello world") == ncodeunits("hello world")
+    @test Base.annotate!(aio, 7:11, :tag => 2) === aio
+    @test Base.annotate!(aio, 1:5, :tag => 1) === aio
+    @test Base.annotations(aio) == [(1:5, :tag => 1), (7:11, :tag => 2)]
     # Reading
     @test read(seekstart(deepcopy(aio.io)), String) == "hello world"
     @test read(seekstart(deepcopy(aio)), String) == "hello world"