Efficient way to map data to legend text color in ggplot2

Question

I'm wondering if there's an efficient way to map data onto legend text color in ggplot2, just like we can do with axis text. Reproducible example follows.

First, let's make a plot:

library(ggplot2)
library(dplyr)

drv_counts <- mutate(mpg,
                     drv = case_when(drv == "r" ~ "rear wheel drive",
                                     drv == "4" ~ "4 wheel drive",
                                     drv == "f" ~ "front wheel drive"),
                     model_drv = interaction(model, drv)) %>%
  group_by(model_drv) %>%
  summarize(model = model[1], drv = drv[1], count = n()) %>%
  arrange(drv, count) %>%
  mutate(model = factor(model, levels = model))

p <- ggplot(drv_counts, aes(x=model, y=count, fill=drv)) +
  geom_col() + coord_flip() + guides(fill = guide_legend(reverse=T)) +
  theme_minimal()
p

Now let's color the axis labels by drive train. This is very easy:

# ggplot2 colors
cols <- c("4 wheel drive" = "#F8766D", "front wheel drive" = "#00BA38", "rear wheel drive" = "#619CFF")

p2 <- p + theme(axis.text.y = element_text(color = cols[drv_counts$drv]))
p2

Now let's try the same trick on the legend. It doesn't work:

p2 + theme(legend.text = element_text(color = cols))

The reason this doesn't work for legend text but does work for axis text is that all the axis labels are drawn in one grob, and hence we can give that grob a vector of colors, but the legend labels are drawn in separate grobs.

We can go in and color all the grobs manually, but that's super ugly and cumbersome:

g <- ggplotGrob(p2)
g$grobs[[15]]$grobs[[1]]$grobs[[9]]$children[[1]]$gp$col <- cols[g$grobs[[15]]$grobs[[1]]$grobs[[9]]$children[[1]]$label]
g$grobs[[15]]$grobs[[1]]$grobs[[10]]$children[[1]]$gp$col <- cols[g$grobs[[15]]$grobs[[1]]$grobs[[10]]$children[[1]]$label]
g$grobs[[15]]$grobs[[1]]$grobs[[11]]$children[[1]]$gp$col <- cols[g$grobs[[15]]$grobs[[1]]$grobs[[11]]$children[[1]]$label]
grid::grid.newpage()
grid::grid.draw(g)

My question is: Can somebody think of a way of getting this effect without having to dig down into the grob tree? I'm Ok with a patch to ggplot2 if it's only a few modified lines. Alternatively, can the digging down into the grob tree be automated so I don't have to access child grobs by manually setting list indices that will change the moment I make a minor change to the figure?

Update: A related question can be found here. To make my question distinct, let's add the requirement that colors aren't copied over from the symbols but rather can be set to any arbitrary values. This added requirement has real-world relevance because I usually use a darker color for text than for symbols.

Answer 1

Here's a pretty mediocre method of hacking grobs together to make a legend. I setup a palette based on the unique values of the drv variable (so it can be scaled to larger datasets or more colors). Then I mapped over the values of the palette to make each legend item: a rectGrob and a textGrob, both with the corresponding color from the palette. These could definitely be tweaked to look better. All of these get arranged into a new grob and stuck alongside the plot with cowplot. It isn't gorgeous but it might be a start.

library(tidyverse)
library(grid)
library(gridExtra)

pal <- colorspace::qualitative_hcl(n = length(unique(drv_counts$drv)), l = 60, c = 70) %>%
  setNames(unique(drv_counts$drv))

p2 <- ggplot(drv_counts, aes(x=model, y=count, fill=drv)) +
  geom_col() + 
  coord_flip() +
  theme_minimal() +
  scale_fill_manual(values = pal, guide = F) +
  theme(axis.text.y = element_text(color = pal[drv_counts$drv]))

legend <- pal %>%
  imap(function(col, grp) {
    rect <- rectGrob(x = 0, width = unit(0.5, "line"), height = unit(0.5, "line"), gp = gpar(col = col, fill = col), hjust = 0)
    label <- textGrob(label = grp, gp = gpar(col = colorspace::darken(col, 0.4), fontsize = 10), x = 0, hjust = 0)
    cowplot::plot_grid(rect, label, nrow = 1, rel_widths = c(0.12, 1))
  }) %>%
  arrangeGrob(grobs = rev(.), padding = unit(0.1, "line"), heights = rep(unit(1.1, "line"), 3))

cowplot::plot_grid(p2, legend, rel_widths = c(1, 0.45))

Created on 2018-05-26 by the reprex package (v0.2.0).