From 921b71ceccdd77f24bfe7a048c642b4c12d33a5c Mon Sep 17 00:00:00 2001
From: Quentin Marcou <18257721+qmarcou@users.noreply.github.com>
Date: Wed, 29 May 2024 14:44:30 +0200
Subject: [PATCH] Fix returned object for transition "all"

---
 R/renewnetTPreg.R | 24 ++++++++++++++++++------
 1 file changed, 18 insertions(+), 6 deletions(-)

diff --git a/R/renewnetTPreg.R b/R/renewnetTPreg.R
index c5a7f33..a4f78a0 100644
--- a/R/renewnetTPreg.R
+++ b/R/renewnetTPreg.R
@@ -321,9 +321,6 @@ renewnetTPreg <- function(s = 0,
       stop("for the transition '23' argument 's' must be larger than 0")
     }
 
-    co <- vector("list", 4)
-    names(co) <- c("co11", "co12", "co13", "co23")
-
     # Check correctness of time_dep_popvars and its interplay with rmap
     rmapsub <- substitute(rmap)
     if (is.list(time_dep_popvars) || is.character(time_dep_popvars)) {
@@ -340,12 +337,16 @@ renewnetTPreg <- function(s = 0,
       stop("`time_dep_popvars` must be a list or vector of strings or NULL")
     }
 
-    if (trans == "all") {
+    if (trans == "all") {  # FIXME name clashes with for loop var name
       transitions <- c("11", "12", "22", "13", "23")
     } else {
       transitions <- c(trans)
     }
 
+    # Pre-allocate list names for the returned object
+    co <- vector("list", length(transitions))
+    names(co) <- glue::glue("co{transitions}")
+
     for (trans in transitions) {
       # Select individuals at risk at time s
       if (trans %in% c("11", "12", "13")) {
@@ -500,8 +501,8 @@ renewnetTPreg <- function(s = 0,
           UPL = extract_stat_matrix("ci.ub"),
           n.failed.boot = NULL
         )
-      if (trans == "all") {
-        co$co11 <- CO
+      if (length(transitions) > 1) {
+        co[[glue::glue("co{trans}")]] <- CO
       } else {
         co <-
           list(
@@ -517,6 +518,17 @@ renewnetTPreg <- function(s = 0,
         return(co)
       }
     }
+    # Get calls and formula for multiple transitions
+    if (length(transitions) > 1) {
+      co$call <- match.call()
+      co$formula <- formula
+      co$transition <- "all"
+      co$s <- s
+      co$t <- t
+      co$n.misobs <- n.misobs
+      class(co) <- "TPreg"
+      return(co)
+    }
   }
 
 estimate_censoring_dist <-