Make lag_forecast() and shift_forecast() return the appropriate classes.

harphub · Jan 22, 2024 · eb2b34b · eb2b34b
1 parent a3d3f19
commit eb2b34b
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 5 deletions.
diff --git a/NAMESPACE b/NAMESPACE
@@ -53,7 +53,7 @@ S3method(scale_point_forecast,default)
 S3method(scale_point_forecast,harp_list)
 S3method(select,harp_fcst)
 S3method(shift_forecast,default)
-S3method(shift_forecast,harp_fcst)
+S3method(shift_forecast,harp_list)
 S3method(spread_members,default)
 S3method(spread_members,harp_fcst)
 S3method(transmute,harp_fcst)

diff --git a/R/lag_forecast.R b/R/lag_forecast.R
@@ -57,7 +57,8 @@ lag_forecast.default <- function(fcst_df, fcst_model, parent_cycles, direction =
     )
 
   purrr::map_dfr(split(fcst_df, fcst_df$parent_cycle), lag_cycle, direction) %>%
-    tidyr::drop_na()
+    tidyr::drop_na() %>%
+    harpCore::as_harp_df()
 
 }
 

diff --git a/R/shift_forecast.R b/R/shift_forecast.R
@@ -45,12 +45,12 @@ shift_forecast.default <- function(.fcst, fcst_shifts, drop_negative_lead_times
     .fcst <- dplyr::filter(.fcst, .data$lead_time >= 0)
   }
 
-  .fcst
+  harpCore::as_harp_df(.fcst)
 
 }
 
 #' @export
-shift_forecast.harp_fcst <- function(.fcst, fcst_shifts, keep_unshifted = FALSE, drop_negative_lead_times = TRUE) {
+shift_forecast.harp_list <- function(.fcst, fcst_shifts, keep_unshifted = FALSE, drop_negative_lead_times = TRUE) {
 
   if (!is.list(fcst_shifts)) {
     if (length(fcst_shifts) > 1) {
@@ -94,6 +94,6 @@ shift_forecast.harp_fcst <- function(.fcst, fcst_shifts, keep_unshifted = FALSE,
       USE.NAMES = FALSE
     )
   }
-  .fcst
+  harpCore::as_harp_list(.fcst)
 }