From a402e0d796b1595ba630e97a6439d8f2debeae2f Mon Sep 17 00:00:00 2001
From: Teun van den Brand <tahvdbrand@gmail.com>
Date: Thu, 18 Jan 2024 15:58:02 +0100
Subject: [PATCH] extract strata more directly in ggplot2 3.5.0

---
 R/add_highlight.R | 37 +++++++++++++++++++++----------------
 1 file changed, 21 insertions(+), 16 deletions(-)

diff --git a/R/add_highlight.R b/R/add_highlight.R
index 4b03b5e8..7460d382 100644
--- a/R/add_highlight.R
+++ b/R/add_highlight.R
@@ -84,26 +84,31 @@ add_highlight.ggsurvfit <- function(gg = NULL,
 
   # Extract names of strata objects
   gg_gb <- ggplot2::ggplot_build(gg)
-  gg_gtable <- ggplot2::ggplot_gtable(gg_gb)
-  gg_guidebox_id <- base::which(base::sapply(
-    gg_gtable$grobs,
-    function(x) x$name
-  ) == "guide-box")
-  gg_table_grob <- gg_gtable$grobs[[gg_guidebox_id]]$grobs[[1]]
 
-  # Get IDs of elements containing strata labels
-  strata_label_ids <- base::grep("label", gg_table_grob$layout$name)
+  if ("get_guide_data" %in% getNamespaceExports("ggplot2")) {
+    strata_labels <- ggplot2::get_guide_data(gg_gb, "colour")$.label
+  } else {
+    gg_gtable <- ggplot2::ggplot_gtable(gg_gb)
+    gg_guidebox_id <- base::which(base::sapply(
+      gg_gtable$grobs,
+      function(x) x$name
+    ) == "guide-box")
+    gg_table_grob <- gg_gtable$grobs[[gg_guidebox_id]]$grobs[[1]]
 
-  extract_strata_name_by_id <- function(gg_table_grob, id) {
-    label <- gg_table_grob$grobs[[id]]$children[[1]]$children[[1]]$label
+    # Get IDs of elements containing strata labels
+    strata_label_ids <- base::grep("label", gg_table_grob$layout$name)
 
-    return(label)
-  }
+    extract_strata_name_by_id <- function(gg_table_grob, id) {
+      label <- gg_table_grob$grobs[[id]]$children[[1]]$children[[1]]$label
+
+      return(label)
+    }
 
-  strata_labels <- base::sapply(strata_label_ids,
-    extract_strata_name_by_id,
-    gg_table_grob = gg_table_grob
-  )
+    strata_labels <- base::sapply(strata_label_ids,
+                                  extract_strata_name_by_id,
+                                  gg_table_grob = gg_table_grob
+    )
+  }
 
   base::sapply(c(strata), function(s) {
     if (!(s %in% strata_labels)) {