change to pull the subset argument

Author: ShreyaSreeram27

R/survfit2_p.R

@@ -42,11 +42,32 @@ survfit2_p <- function(x, pvalue_fun = format_p, prepend_p = TRUE, rho = 0) {
     cli_abort("The {.fun survfit2_p} does not support multi-state models.")
   }
 
-  survival::survdiff(
-    formula = .extract_formula_from_survfit(x),
-    data = .extract_data_from_survfit(x),
-    rho = rho
-  ) %>%
+  #extract subset argument from the original survfit call 
+  call_list <- x$call %>% as.list()
+  subset_arg <- call_list[["subset"]]
+  
+  #call survdiff with subset argument if it exists
+  if (!is.null(subset_arg)) {
+    # do.call to properly pass the subset argument
+    survdiff_result <- do.call(
+      survival::survdiff,
+      list(
+        formula = .extract_formula_from_survfit(x),
+        data = .extract_data_from_survfit(x),
+        subset = subset_arg,
+        rho = rho
+      ),
+      envir = x$.Environment
+    )
+  } else {
+    survdiff_result <- survival::survdiff(
+      formula = .extract_formula_from_survfit(x),
+      data = .extract_data_from_survfit(x),
+      rho = rho
+    )
+  }
+  
+  survdiff_result %>%
     broom::glance() %>%
     dplyr::pull("p.value") %>%
     pvalue_fun() %>%
@@ -55,4 +76,4 @@ survfit2_p <- function(x, pvalue_fun = format_p, prepend_p = TRUE, rho = 0) {
       prepend_p & grepl(pattern = "^<|^>", x = .) ~ paste0("p", .),
       prepend_p ~ paste0("p=", .)
     )}
-}
+}