add specialisation for vector of vectors (#5103)

Co-authored-by: t-bltg <tf.bltg@gmail.com>

Author: InfiniteChai

RecipesPipeline/src/type_recipe.jl

@@ -54,6 +54,52 @@ function _apply_type_recipe(plotattributes, v::AbstractArray, letter)
     w
 end
 
+# Specialisation to apply type recipes on a vector of vectors. The type recipe can either 
+# apply to the vector of elements or the elements themselves
+function _apply_type_recipe(plotattributes, v::AVec{<:AVec}, letter)
+    plt = plotattributes[:plot_object]
+    preprocess_axis_attrs!(plt, plotattributes, letter)
+    # First we attempt on the vector of vector type recipe across everything. 
+    w = RecipesBase.apply_recipe(plotattributes, typeof(v), v)[1].args[1]
+    warn_on_recipe_aliases!(plt, plotattributes, :type, v)
+    if typeof(v) != typeof(w)
+        postprocess_axis_attrs!(plt, plotattributes, letter)
+        return w
+    end
+    # Next we attempt the array type recipe and if any of the vector elements applies,
+    # we will stop there. Note we use the same type equivalency test as for a general array 
+    # to check if changes applied
+    did_replace = false
+    w = map(v) do u
+        newu = RecipesBase.apply_recipe(plotattributes, typeof(u), u)[1].args[1]
+        warn_on_recipe_aliases!(plt, plotattributes, :type, u)
+        did_replace |= typeof(u) !== typeof(newu)
+        newu
+    end
+
+    # if nothing changed, then we attempt it at a piecewise level
+    if !did_replace
+        if (smv = skipmissing(Base.Iterators.flatten(v))) |> isempty
+            postprocess_axis_attrs!(plt, plotattributes, letter)
+            # We'll just leave it untampered with if there are no elements
+            return v
+        end
+        x = first(smv)
+        args = RecipesBase.apply_recipe(plotattributes, typeof(x), x)[1].args
+        warn_on_recipe_aliases!(plt, plotattributes, :type, x)
+        postprocess_axis_attrs!(plt, plotattributes, letter)
+        return if length(args) == 2 && all(arg -> arg isa Function, args)
+            numfunc, formatter = args
+            Formatted(map(u -> map(numfunc, u), v), formatter)
+        else
+            v
+        end
+    end
+
+    postprocess_axis_attrs!(plt, plotattributes, letter)
+    w
+end
+
 # special handling for Surface... need to properly unwrap and re-wrap
 _apply_type_recipe(
     plotattributes,

RecipesPipeline/test/Project.toml

@@ -1,5 +1,6 @@
 [deps]
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
 RecipesPipeline = "01d81517-befc-4cb6-b9ec-a95719d0359c"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"

RecipesPipeline/test/runtests.jl

@@ -3,8 +3,10 @@ using BenchmarkTools
 using StableRNGs
 using Test
 
-import RecipesPipeline: _prepare_series_data
+import Dates
+import RecipesPipeline: _prepare_series_data, _apply_type_recipe
 import RecipesBase
+import RecipesBase: @recipe
 
 @testset "DefaultsDict" begin
     dd = DefaultsDict(Dict(:foo => 1, :bar => missing), Dict(:foo => nothing, :baz => 'x'))
@@ -82,6 +84,41 @@ end
     @test RecipesBase.is_key_supported("key")
 end
 
+@testset "_apply_type_recipe" begin
+    plt = nothing
+    plotattributes = Dict{Symbol,Any}(:plot_object => plt)
+    @test _apply_type_recipe(plotattributes, [1, 2, 3], :x) == [1, 2, 3]
+    @test _apply_type_recipe(plotattributes, [[1, 2], [3, 4]], :x) == [[1, 2], [3, 4]]
+    res = _apply_type_recipe(plotattributes, [Dates.Date(2001)], :x)
+    @test typeof(res) <: Formatted
+    @test res.data == [Dates.value(Dates.Date(2001))]
+    @test res.formatter(Dates.value(Dates.Date(2001))) == "2001-01-01"
+
+    res = _apply_type_recipe(plotattributes, [[Dates.Date(2001)]], :x)
+    @test typeof(res) <: Formatted
+    @test res.data == [[Dates.value(Dates.Date(2001))]]
+    @test res.formatter(Dates.value(Dates.Date(2001))) == "2001-01-01"
+
+    struct Test1 <: Number
+        val::Float64
+    end
+
+    @recipe f(::Type{T}, v::T) where {T<:AbstractVector{<:Test1}} = map(x -> x.val + 1, v)
+
+    @test _apply_type_recipe(plotattributes, Test1.([1, 2, 3]), :x) == [2.0, 3.0, 4.0]
+    @test _apply_type_recipe(plotattributes, [Test1.([1, 2, 3])], :x) == [[2.0, 3.0, 4.0]]
+
+    struct Test2 <: Number
+        val::Float64
+    end
+
+    @recipe f(::Type{T}, v::T) where {T<:AbstractVector{<:AbstractVector{<:Test2}}} =
+        map(x -> map(y -> y.val + 2, x), v)
+
+    @test _apply_type_recipe(plotattributes, Test2.([1, 2, 3]), :x) == Test2.([1, 2, 3])
+    @test _apply_type_recipe(plotattributes, [Test2.([1, 2, 3])], :x) == [[3.0, 4.0, 5.0]]
+end
+
 @testset "_prepare_series_data" begin
     @test_throws ErrorException _prepare_series_data(:test)
     @test _prepare_series_data(nothing) ≡ nothing