added classprob prediction type

jonthegeek · Jan 29, 2019 · 0ca0445 · 0ca0445
1 parent 4a37edb
commit 0ca0445
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 2 deletions.
diff --git a/R/nullmodel_data.R b/R/nullmodel_data.R
@@ -23,14 +23,28 @@ null_model_parsnip_data <-
       defaults = list()
       ),
     class = list(
+      pre = NULL,
+      post = NULL,
+      func = c(fun = "predict"),
+      args =
+        list(
+          object = quote(object$fit),
+          new_data = quote(new_data),
+          type = "class"
+        )
+    ),
+    classprob = list(
         pre = NULL,
-        post = NULL,
+        post = function(x, object) {
+          str(as_tibble(x))
+          as_tibble(x)
+        },
         func = c(fun = "predict"),
         args =
           list(
             object = quote(object$fit),
             new_data = quote(new_data),
-            type = "class"
+            type = "prob"
           )
       ),
     numeric = list(

diff --git a/tests/testthat/test-nullmodel.R → tests/testthat/test_nullmodel.R b/tests/testthat/test-nullmodel.R → tests/testthat/test_nullmodel.R
@@ -134,3 +134,4 @@ test_that('classification', {
   expect_true(!is.null(null_model$fit))
 })
 
+
Original file line number	Diff line number	Diff line change
Expand Up		@@ -134,3 +134,4 @@ test_that('classification', {
		expect_true(!is.null(null_model$fit))
		})