From 950ba75c8898ca97f933c82b3d31e28c00999ab8 Mon Sep 17 00:00:00 2001
From: Elias Carvalho <73039601+eliascarv@users.noreply.github.com>
Date: Mon, 8 Apr 2024 11:34:17 -0300
Subject: [PATCH] Update color mapping for categorical values (#17)

---
 ext/ColorfyCategoricalArraysExt.jl | 8 ++++----
 test/runtests.jl                   | 9 ++++++++-
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/ext/ColorfyCategoricalArraysExt.jl b/ext/ColorfyCategoricalArraysExt.jl
index 38d4bbb..ca6e7ee 100644
--- a/ext/ColorfyCategoricalArraysExt.jl
+++ b/ext/ColorfyCategoricalArraysExt.jl
@@ -7,14 +7,14 @@ module ColorfyCategoricalArraysExt
 using Colorfy
 using Colorfy: Values
 using ColorSchemes: colorschemes
-using CategoricalArrays: CategoricalValue, levelcode
+using CategoricalArrays: CategoricalValue, levels, levelcode
 
 function Colorfy.getcolors(colorfier::Colorfier{<:Values{CategoricalValue}})
   values = Colorfy.values(colorfier)
   colorscheme = Colorfy.colorscheme(colorfier)
-  colorscheme[levelcode.(values)]
+  nlevels = length(levels(values))
+  categcolors = colorscheme[range(0, 1, length=nlevels)]
+  categcolors[levelcode.(values)]
 end
 
-Colorfy.defaultcolorscheme(::Values{CategoricalValue}) = colorschemes[:Set3_9]
-
 end
diff --git a/test/runtests.jl b/test/runtests.jl
index 3893bea..96041c9 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -107,7 +107,14 @@ using Test
 
   @testset "CategoricalArrays" begin
     values = categorical(["n", "n", "y", "y", "n", "y"], levels=["y", "n"])
-    colors = colorschemes[:Set3_9][[2, 2, 1, 1, 2, 1]]
+    categcolors = colorschemes[:viridis][range(0, 1, length=2)]
+    colors = categcolors[[2, 2, 1, 1, 2, 1]]
+    @test colorfy(values) == coloralpha.(colors, 1)
+    @test colorfy(values, alphas=0.5) == coloralpha.(colors, 0.5)
+
+    values = categorical([2, 1, 1, 3, 1, 3, 3, 2, 1, 2], levels=1:3)
+    categcolors = colorschemes[:viridis][range(0, 1, length=3)]
+    colors = categcolors[[2, 1, 1, 3, 1, 3, 3, 2, 1, 2]]
     @test colorfy(values) == coloralpha.(colors, 1)
     @test colorfy(values, alphas=0.5) == coloralpha.(colors, 0.5)
   end