Fix subsetting that failed on single column df. Close tidymodels#119

antoine-sachet · Oct 28, 2020 · 170e571 · 170e571
1 parent ab056fc
commit 170e571
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 5 deletions.
diff --git a/R/cor_df.R b/R/cor_df.R
@@ -5,7 +5,7 @@ as_matrix.cor_df <- function(x, diagonal) {
 
   # Separate rownames
   row_name <- x$rowname
-  x <- x[, colnames(x) != "rowname"]
+  x <- x[colnames(x) != "rowname"]
   # Convert to matrix and set rownames
   class(x) <- "data.frame"
   x <- as.matrix(x)
@@ -22,7 +22,7 @@ shave.cor_df <- function(x, upper = TRUE) {
 
   # Separate rownames
   row_name <- x$rowname
-  x <- x[, colnames(x) != "rowname"]
+  x <- x[colnames(x) != "rowname"]
 
   # Remove upper matrix
   if (upper) {
@@ -93,7 +93,7 @@ focus_if.cor_df <- function(x, .predicate, ..., mirror = FALSE) {
 
   # Identify which variables to keep
   to_keep <- map_lgl(
-    x[, colnames(x) != "rowname"],
+    x[colnames(x) != "rowname"],
     .predicate, ...
   )
 

diff --git a/R/reshape.R b/R/reshape.R
@@ -124,7 +124,7 @@ stretch <- function(x, na.rm = FALSE, remove.dups =  FALSE) {
 stretch.cor_df <- function(x, na.rm = FALSE, remove.dups =  FALSE) {
   if(remove.dups) x <- shave(x)
   row_name <- x$rowname
-  x <- x[, colnames(x) != "rowname"]
+  x <- x[colnames(x) != "rowname"]
   tb <- imap_dfr(x, ~tibble(x = .y, y = row_name, r = .x))
   if(na.rm) tb <- tb[!is.na(tb$r), ]
   if(remove.dups) {

diff --git a/R/utility.R b/R/utility.R
@@ -20,7 +20,7 @@ as_cordf <- function(x, diagonal = NA) {
   }
   x <- as.data.frame(x)
   row_name <- x$rowname
-  x <- x[, colnames(x) != "rowname"]
+  x <- x[colnames(x) != "rowname"]
   rownames(x) <- row_name
   if(ncol(x) != nrow(x)) {
     stop("Input object x is not a square. ",

diff --git a/tests/testthat/test-as_cordf.R b/tests/testthat/test-as_cordf.R
@@ -18,3 +18,9 @@ test_that("Diagonal sets correctly", {
   expect_equal(all(is.na(diag(as.matrix(as_cordf(d, diagonal = NA)[, -1])))), TRUE)
   expect_equal(all(diag(as.matrix(as_cordf(d, diagonal = 100)[, -1] == 100))), TRUE)
 })
+
+test_that("as_cordf handles single correlation", {
+  d1 <- cor(mtcars["cyl"])
+  expect_s3_class(as_cordf(d1), "cor_df")
+  expect_equal(colnames(as_cordf(d1)), c("rowname", colnames(d1)))
+})