From adff5c1dc1a3bf0e63832022f3bb4465494c1e9a Mon Sep 17 00:00:00 2001
From: Michael Mayer <mayermichael79@gmail.com>
Date: Sat, 6 Jul 2024 10:21:48 +0200
Subject: [PATCH] Add checks to prediction output

---
 NEWS.md                     | 3 ++-
 R/utils.R                   | 8 +++++++-
 tests/testthat/test-utils.R | 5 +++++
 3 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/NEWS.md b/NEWS.md
index e54a340..05989c2 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -4,12 +4,13 @@ This release is intended to be the last before stable version 1.0.0.
 
 ## Major changes
 
-- Factor-valued predictions are not anymore supported.
+- Factor-valued predictions are not supported anymore.
 
 ## Maintenance
 
 - Fix CRAN note about unavailable link to `gam::gam()`.
 - Added dependency to {MASS} for calculating Moore-Penrose generalized matrix inverse.
+- Slight code improvements.
 
 # kernelshap 0.5.0
 
diff --git a/R/utils.R b/R/utils.R
index d1c6ee6..59268c1 100644
--- a/R/utils.R
+++ b/R/utils.R
@@ -153,7 +153,13 @@ align_pred <- function(x) {
   if (is.data.frame(x) && ncol(x) == 1L) {
     x <- x[[1L]]
   }
-  as.matrix(x)
+  if (!is.matrix(x)) {
+    x <- as.matrix(x)
+  }
+  if (!is.numeric(x) && !is.logical(x)) {
+    stop("Predictions must be numeric!")
+  }
+  return(x)
 }
 
 #' Head of List Elements
diff --git a/tests/testthat/test-utils.R b/tests/testthat/test-utils.R
index 7d0ea9b..1134de7 100644
--- a/tests/testthat/test-utils.R
+++ b/tests/testthat/test-utils.R
@@ -82,3 +82,8 @@ test_that("wrowmean_vector() works for 1D matrices", {
   expect_equal(out2, cbind(a = c(a, b)))
 })
 
+test_that("align_pred() works", {
+  expect_error(align_pred(c("A", "B")))
+  expect_error(align_pred(factor(c("A", "B"))))
+  expect_equal(align_pred(1:4), as.matrix(1:4))
+})