refac: function for pulling heatmap genotype data from ggtree object

R/plot_tree.R

@@ -197,6 +197,32 @@ append_interactivity_data <- function(ggobj,
   ggobj
 }
 
+#' Extract the data about viral genotypes from a \code{ggtree} object
+#'
+#' @param   ggobj   A \code{ggtree} object, as generated by \code{append_interactivity_data}. The
+#'   \code{data} entry for this object should contain the columns "node", "label" and a column for
+#'   each of the \code{mut_regex} values.
+#' @param   n_leaves   Scalar integer. The number of leaves in the tree.
+#' @param   mut_regex   String. Regular expression defining the mutations under study here. This
+#'   should be a subset of the column-names in \code{ggobj$data}.
+#'
+#' @return   Data-frame.
+
+extract_genotype_data <- function(ggobj,
+                                  n_leaves,
+                                  mut_regex) {
+  genotype <- as.data.frame(
+    ggobj$data[
+      ggobj$data$node <= n_leaves,
+      c("label", mut_regex)
+    ]
+  )
+  rownames(genotype) <- genotype$label
+  genotype <- genotype[, -1, drop = FALSE]
+
+  genotype
+}
+
 #' Adds a heatmap to the right of a ggtree object
 #'
 #' @param   ggobj   A ggtree object.

R/treeview.R

@@ -254,14 +254,11 @@ treeview <- function(e0,
       mut_regex = mut_regex
     )
 
-    genotype <- as.data.frame(
-      gtr1.1$data[
-        gtr1.1$data$node <= ape::Ntip(tr2),
-        c("label", mut_regex)
-      ]
+    genotype <- extract_genotype_data(
+      ggobj = gtr1.1,
+      n_leaves = ape::Ntip(tr2),
+      mut_regex = mut_regex
     )
-    rownames(genotype) <- genotype$label
-    genotype <- genotype[, -1, drop = FALSE]
 
     gtr1.2 <- append_heatmap(
       ggobj = gtr1.1,