Add shaded confidence interval to ggflexsurv plot

Question

ggflexsurv from the package survminer seems to only support dashed confidence interval lines rather than shaded/filled. I've tried a few work arounds, such as this thread but to no avail.

My last attempt was to inset the code from the answer in the linked thread into the source code for ggflexsurv

ggflexsurvplotmod <- function(fit, 
                               data = NULL,
                               fun = c("survival", "cumhaz"),
                               summary.flexsurv = NULL,
                               size = 1, 
                               conf.int = FALSE,
                               conf.int.flex = conf.int, 
                               conf.int.km = FALSE,
                               legend.labs = NULL,...)
  
{
  
  if (!requireNamespace("flexsurv", quietly = TRUE)){
    stop("flexsurv package needed for this function to work. Please install it.")}
  
  if(!inherits(fit, "flexsurvreg")){
    stop("Can't handle an object of class ", class(fit))
    fun <- match.arg(fun)}
  
  data <- .get_data(fit, data = data, complain = FALSE)
  
  summ <- .summary_flexsurv(fit, type = fun,
                            summary.flexsurv = summary.flexsurv)
  
  .strata <- summ$strata
  
  n.strata <- .strata %>% .levels() %>% length()
  
  fit.ext <- .extract.survfit(fit)
  
  surv.obj <- fit.ext$surv
  
  surv.vars <- fit.ext$variables
  
  .formula <- fit.ext$formula
  
  isfac <- .is_all_covariate_factor(fit)
  
  if(!all(isfac))
  {.formula <- .build_formula(surv.obj, "1")
    n.strata <- 1}
  
  if(n.strata == 1 & missing(conf.int)) conf.int <- TRUE
  
  # Fit KM survival curves
  
  #::::::::::::::::::::::::::::::::::::::
x <- do.call(survival::survfit, 
             list(formula = .formula, data = data))

fun <- if(fun == "survival") NULL else fun

ggsurv <- ggsurvplot_core(x, 
                          data = data, 
                          size = 0.5,
                          fun = fun, 
                          conf.int = conf.int.km,
                          legend.labs = legend.labs, ...)
  
  # Overlay the fitted models
  
  #::::::::::::::::::::::::::::::::::::::
  
  # Check legend labels if specified
  
  if(!is.null(legend.labs)){
    
    if(n.strata != length(legend.labs))
      stop("The length of legend.labs should be ", n.strata )
    
    summ$strata <- factor(summ$strata,
                          levels = .levels(.strata),
                          labels = legend.labs)
    
  }
  
  time <- est <- strata <- lcl <- ucl <- NULL

ggsurv$plot <- 
    ggsurv$plot +
    geom_line(aes(time, 
                  est, 
                  color = strata),
              data = summ, 
              size = size)
  
  #CODE FROM PREVIOUS STACKOVERFLOW THREAD#######
  
  stairstepn <- function( data, direction="hv", yvars="y" )
    {direction <- match.arg(direction, c("hv", "vh") )
    data <- as.data.frame(data)[ order( data$x ), ]
    n <- nrow( data )
  
  if ( direction == "vh" ) {
    xs <- rep( 1:n, each = 2 )[ -2 * n ]
    ys <- c( 1, rep( 2:n, each = 2 ) )
  } else {
    ys <- rep( 1:n, each = 2 )[ -2 * n ]
    xs <- c( 1, rep( 2:n, each = 2))
  }
  
  data.frame(
    x = data$x[ xs ]
    , data[ ys, yvars, drop=FALSE ]
    , data[ xs, setdiff( names( data ), c( "x", yvars ) ), drop=FALSE ]
  ) 
  
  }
  
  stat_stepribbon <- function(mapping = NULL, 
                              data = NULL, 
                              geom = "ribbon", 
                              position = "identity", 
                              inherit.aes = TRUE) 
  {ggplot2::layer(stat = Stepribbon, 
                    mapping = mapping, 
                  data = data, 
                  geom = geom,
                  position = position, 
                  inherit.aes = inherit.aes)}
  
  StatStepribbon <- ggproto("stepribbon", 
                            Stat,
                            compute_group = 
                              function(., data, scales, 
                                       direction = "hv",
                                       yvars = c( "ymin", "ymax" ), ...)
                                {stairstepn(data = data,
                                            direction = direction,
                                            yvars = yvars )},
                            required_aes = c( "x", "ymin", "ymax" ))
  
  ###########################################################################
  
  if(conf.int.flex) ggsurv$plot <- ggsurv$plot+
    geom_ribbon(aes(ymin = lcl,
                    ymax = ucl), 
                fill = strata, 
                stat = "stepribbon", 
                data = summ, 
                alpha=0.3)
  
  #other version of geom_ribbon() attempted
  #geom_ribbon(aes(ymin=lcl,ymax=ucl), data = summ,fill=strata, alpha=0.3) #stat="identity"
  #original code in package
  #geom_line(aes(time, lcl, color = strata), data = summ, #size = 0.5, linetype = "dashed")+
  
  #geom_line(aes(time, ucl, color = strata), data = summ, #size = 0.5, linetype = "dashed")
  
  ggsurv
  
}

.summary_flexsurv <- function(fit, type = "survival", summary.flexsurv = NULL){
  
  summ <- summary.flexsurv
  
  if(is.null(summary.flexsurv)){
    summ <- summary(fit, type = type)}
  
  if(length(summ) == 1){summ <- summary(fit)[[1]] %>%
    dplyr::mutate(strata = "All")}
  
  else{.strata <- names(summ)
  summ <- purrr::pmap(list(.strata, summ), function(.s, .summ){
    dplyr::mutate(.summ, strata = .s )})
  summ <- dplyr::bind_rows(summ)
  summ$strata <- factor(summ$strata, levels = .strata)}
  
  summ
  
}

# Check if all covariates are factor or character vector

.is_all_covariate_factor <- function(fit){
  x <- fit
mf <- stats::model.frame(x)
Xraw <- mf[,attr(mf, "covnames.orig"), drop=FALSE]
dat <- x$datasapply(Xraw,is_factor_or_character)}

is_factor_or_character <- function(x){is.facet(x) | is.character(x)}

#This code "saves" the modified code into the packageenvironment(ggflexsurvplotmod) <- asNamespace('survminer')
assignInNamespace("ggflexsurvplot", ggflexsurvplotmod, ns = "survminer")

I then run the following

 code_example =  flexsurvreg(formula = Surv(time, status) ~ sex, data = lung,dist = "gompertz")

ggflexsurvplotmod(code_example, conf.int = T) `

which results in the following error

Error in geom_ribbon(): Problem while computing aesthetics. i Error occurred in the 5th layer. Caused by error: object 'surv' not found`