Merge pull request #478 from EmilHvitfeldt/augment-tibble

juliasilge · web-flow · commit c7e1010f45ac · 2021-04-26T09:28:03.000-06:00
make augment() return tibbles
diff --git a/R/augment.R b/R/augment.R
@@ -77,5 +77,5 @@ augment.model_fit <- function(x, new_data, ...) {
   } else {
     rlang::abort(paste("Unknown mode:", x$spec$mode))
   }
-  new_data
+  as_tibble(new_data)
 }
diff --git a/tests/testthat/test-augment.R b/tests/testthat/test-augment.R
@@ -35,6 +35,8 @@ test_that('regression models', {
   )
   expect_equal(nrow(augment(reg_xy, head(mtcars[, -1]))), 6)
 
+  expect_s3_class(augment(reg_form, head(mtcars)), "tbl_df")
+
   reg_form$spec$mode <- "depeche"
 
   expect_error(augment(reg_form, head(mtcars[, -1])), "Unknown mode: depeche")

Original file line number	Diff line number	Diff line change
`@@ -77,5 +77,5 @@ augment.model_fit <- function(x, new_data, ...) {`
`77`	`77`	`} else {`
`78`	`78`	`rlang::abort(paste("Unknown mode:", x$spec$mode))`
`79`	`79`	`}`
`80`		`- new_data`
	`80`	`+ as_tibble(new_data)`
`81`	`81`	`}`
Original file line number	Diff line number	Diff line change
`@@ -35,6 +35,8 @@ test_that('regression models', {`
`35`	`35`	`)`
`36`	`36`	`expect_equal(nrow(augment(reg_xy, head(mtcars[, -1]))), 6)`
`37`	`37`
	`38`	`+ expect_s3_class(augment(reg_form, head(mtcars)), "tbl_df")`
	`39`	`+`
`38`	`40`	`reg_form$spec$mode <- "depeche"`
`39`	`41`
`40`	`42`	`expect_error(augment(reg_form, head(mtcars[, -1])), "Unknown mode: depeche")`