Merge pull request #2069 from rstudio/fix-stub-helper

fix: ensure `stub()` helper works with multi-column stubs

Author: rich-iannone

R/dt_stub_df.R

@@ -74,6 +74,24 @@ dt_stub_df_init <- function(
         data <- dt_boxhead_set_stub(data = data, var = col)
       }
 
+      # Reorder the boxhead to match the order specified in rowname_col
+      # This ensures stub columns appear in the order the user specified
+      dt_boxhead <- dt_boxhead_get(data = data)
+      
+      # Split boxhead into stub and non-stub rows
+      stub_rows <- dt_boxhead[dt_boxhead$var %in% rowname_col, ]
+      non_stub_rows <- dt_boxhead[!dt_boxhead$var %in% rowname_col, ]
+      
+      # Reorder stub rows to match rowname_col order
+      stub_rows <- stub_rows[match(rowname_col, stub_rows$var), ]
+      
+      # Combine: stub rows first (in specified order), then non-stub rows
+      dt_boxhead <- rbind(stub_rows, non_stub_rows)
+      rownames(dt_boxhead) <- NULL
+      
+      # Update the boxhead
+      data <- dt_boxhead_set(data = data, boxh = dt_boxhead)
+      
       # Use the rightmost column as the primary row ID
       rightmost_col <- rowname_col[length(rowname_col)]
       rownames <- data_tbl[[rightmost_col]]

R/helpers.R

@@ -999,6 +999,14 @@ adjust_luminance <- function(
 #' operations, the `stub()` select helper can be used. This obviates the need
 #' to use the name of the column that was selected as the stub column.
 #'
+#' @param column *Stub column selection*
+#'
+#'   `scalar<integer>` // *default:* `NULL` (optional)
+#'
+#'   For tables with multi-column stubs, optionally specify which stub column to
+#'   target. Use `1` for the rightmost stub column, `2` for the second from
+#'   right, etc. If `NULL` (the default), all stub columns are selected.
+#'
 #' @return A character vector of class `"stub_column"`.
 #'
 #' @section Examples:
@@ -1031,6 +1039,40 @@ adjust_luminance <- function(
 #' `r man_get_image_tag(file = "man_stub_1.png")`
 #' }}
 #'
+#' For multi-column stubs, you can set widths for all stub columns together or
+#' target specific columns individually:
+#'
+#' ```r
+#' exibble |>
+#'   dplyr::select(num, char, row, group) |>
+#'   gt(rowname_col = c("group", "row")) |>
+#'   cols_width(
+#'     stub() ~ px(200),        # All stub columns get 200px
+#'     everything() ~ px(100)
+#'   )
+#' ```
+#'
+#' \if{html}{\out{
+#' `r man_get_image_tag(file = "man_stub_2.png")`
+#' }}
+#'
+#' Or target specific stub columns (1 = rightmost, 2 = second from right):
+#'
+#' ```r
+#' exibble |>
+#'   dplyr::select(num, char, row, group) |>
+#'   gt(rowname_col = c("group", "row")) |>
+#'   cols_width(
+#'     stub(1) ~ px(70),        # Rightmost stub column (row)
+#'     stub(2) ~ px(200),       # Second stub column (group)
+#'     everything() ~ px(100)
+#'   )
+#' ```
+#'
+#' \if{html}{\out{
+#' `r man_get_image_tag(file = "man_stub_3.png")`
+#' }}
+#'
 #' @family helper functions
 #' @section Function ID:
 #' 8-10
@@ -1039,9 +1081,10 @@ adjust_luminance <- function(
 #' `v0.8.0` (November 16, 2022)
 #'
 #' @export
-stub <- function() {
+stub <- function(column = NULL) {
   x <- "::stub::"
   class(x) <- "stub_column"
+  attr(x, "column") <- column
   x
 }
 

R/resolver.R

@@ -429,14 +429,48 @@ resolve_cols_i <- function(
     # If we use the gt-specific select helper `stub()` then we
     # will retrieve the stub var name and return the output in the
     # same format as the return value for `tidyselect::eval_select()`
-    if (rlang::as_label(quo) %in% c("stub()", "gt::stub()")) {
+    # Check if the expression is a stub() call (with or without arguments)
+    stub_label <- rlang::as_label(quo)
+    if (grepl("^(gt::)?stub\\(", stub_label)) {
 
       stub_var <- dt_boxhead_get_var_stub(data = data)
 
-      if (!is.null(stub_var)) {
-        stub_col <- 1
-        names(stub_col) <- stub_var
-        return(stub_col)
+      if (!is.null(stub_var) && !all(is.na(stub_var))) {
+        
+        # Evaluate the stub() expression to get the column parameter if any
+        stub_result <- rlang::eval_tidy(quo)
+        column_idx <- attr(stub_result, "column")
+        
+        # If column index is specified, select that specific stub column
+        # (1 = rightmost, 2 = second from right, etc.)
+        if (!is.null(column_idx)) {
+          # Validate the column index
+          if (!is.numeric(column_idx) || length(column_idx) != 1 || column_idx < 1) {
+            cli::cli_abort(
+              "The {.arg column} argument in {.fn stub} must be a single positive integer."
+            )
+          }
+          
+          # Check if the requested column exists
+          if (column_idx > length(stub_var)) {
+            cli::cli_abort(c(
+              "Cannot select stub column {column_idx}.",
+              "i" = "This table has {length(stub_var)} stub column{?s}."
+            ))
+          }
+          
+          # Select from right to left (1 = rightmost)
+          selected_stub_var <- stub_var[length(stub_var) - column_idx + 1]
+          stub_col <- 1
+          names(stub_col) <- selected_stub_var
+          return(stub_col)
+          
+        } else {
+          # No column specified, return all stub columns
+          stub_col <- seq_along(stub_var)
+          names(stub_col) <- stub_var
+          return(stub_col)
+        }
       } else {
         return(NULL)
       }

R/utils_render_html.R

@@ -2492,15 +2492,21 @@ calculate_hierarchical_stub_rowspans <- function(data) {
       # Check if current row should continue the span from previous row
       should_continue_span <- TRUE
 
-      # Must match current column value
-      if (col_values[row_idx] != col_values[row_idx - 1]) {
+      # Must match current column value (handle NAs properly)
+      curr_val <- col_values[row_idx]
+      prev_val <- col_values[row_idx - 1]
+      
+      # Two values match if they're both identical or both NA
+      if (!identical(curr_val, prev_val)) {
         should_continue_span <- FALSE
       }
 
       # Must match all values in columns to the left
       if (col_idx > 1) {
         for (left_col_idx in 1:(col_idx - 1)) {
-          if (stub_matrix[row_idx, left_col_idx] != stub_matrix[row_idx - 1, left_col_idx]) {
+          curr_left <- stub_matrix[row_idx, left_col_idx]
+          prev_left <- stub_matrix[row_idx - 1, left_col_idx]
+          if (!identical(curr_left, prev_left)) {
             should_continue_span <- FALSE
             break
           }

R/utils_render_latex.R

@@ -776,15 +776,20 @@ create_body_component_l <- function(data, colwidth_df) {
         for (row_idx in 2:n_rows) {
           should_hide <- TRUE
 
-          # Check if current value matches previous value
-          if (stub_matrix[row_idx, col_idx] != stub_matrix[row_idx - 1, col_idx]) {
+          # Check if current value matches previous value (handle NAs properly)
+          curr_val <- stub_matrix[row_idx, col_idx]
+          prev_val <- stub_matrix[row_idx - 1, col_idx]
+          
+          if (!identical(curr_val, prev_val)) {
             should_hide <- FALSE
           }
 
           # Also check that all columns to the left match
           if (should_hide && col_idx > 1) {
             for (left_col_idx in 1:(col_idx - 1)) {
-              if (stub_matrix[row_idx, left_col_idx] != stub_matrix[row_idx - 1, left_col_idx]) {
+              curr_left <- stub_matrix[row_idx, left_col_idx]
+              prev_left <- stub_matrix[row_idx - 1, left_col_idx]
+              if (!identical(curr_left, prev_left)) {
                 should_hide <- FALSE
                 break
               }

images/man_stub_2.png

@@ -247,8 +247,9 @@ test_that("The stub gets its alignment set properly with `cols_align()`", {
     c("stub", "vals")
   )
 
-  # Expect an error if using `stub()` when there is no stub
-  expect_error(
+  # When using `stub()` on a table with no stub, it should not error
+  # (consistent with tidyselect behavior where selecting nothing is valid)
+  expect_no_error(
     gt(tbl) |>
     cols_align(align = "center", columns = stub())
   )