add_totals_row now works on two-column input

closes #69
sfirke · Oct 10, 2016 · 5b38697 · 5b38697
1 parent 1cb6eee
commit 5b38697
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 10 deletions.
diff --git a/R/add_totals.R b/R/add_totals.R
@@ -17,7 +17,7 @@
 add_totals_row <- function(dat, na.rm = TRUE){
   check_all_cols_after_first_are_numeric(dat)
   dat[[1]] <- as.character(dat[[1]]) # for binding to the "Total" character value of add-on row
-  col_totals <- data.frame(x1 = "Total", t(colSums(dat[,-1], na.rm = na.rm)), stringsAsFactors = FALSE) %>%
+  col_totals <- data.frame(x1 = "Total", t(colSums(dat[-1], na.rm = na.rm)), stringsAsFactors = FALSE) %>%
     stats::setNames(names(dat))
   dplyr::bind_rows(dat, col_totals)
 }
@@ -39,7 +39,7 @@ add_totals_row <- function(dat, na.rm = TRUE){
 
 add_totals_col <- function(dat, na.rm = TRUE){
   check_all_cols_after_first_are_numeric(dat)
-  row_totals <- data.frame(Total = rowSums(dat[,-1], na.rm = na.rm))
+  row_totals <- data.frame(Total = rowSums(dat[-1], na.rm = na.rm))
   dplyr::bind_cols(dat, row_totals)
 }
 
diff --git a/tests/testthat/test-add-totals.R b/tests/testthat/test-add-totals.R
@@ -4,24 +4,24 @@ library(janitor)
 context("add_totals functions")
 
 library(dplyr)
-dat <- data.frame(a = c(1:3, 1:3, 1, 1, 1),
-                  b = c(rep(c("big", "small", "big"), 3))
+dat <- data.frame(a = c(rep(c("big", "small", "big"), 3)),
+                  b = c(1:3, 1:3, 1, 1, 1)
 )
 ct <- dat %>%
-  crosstab(b, a)
+  crosstab(a, b)
 
 
 test_that("error thrown if column beyond the first is not numeric", {
-  expect_error(add_totals_row(dat),
+  expect_error(add_totals_row(dat %>% select(b, a)),
                "all columns after the first one must be numeric")
-  expect_error(add_totals_col(dat),
+  expect_error(add_totals_col(dat %>% select(b, a)),
                "all columns after the first one must be numeric")
 })
 
 
 test_that("totals row is correct", {
   expect_equal(add_totals_row(ct),
-               data.frame(b = c("big", "small", "Total"),
+               data.frame(a = c("big", "small", "Total"),
                           `1` = c(4, 1, 5),
                           `2` = c(0, 2, 2),
                           `3` = c(2, 0, 2),
@@ -33,7 +33,7 @@ test_that("totals row is correct", {
 
 test_that("totals col is correct", {
   expect_equal(add_totals_col(ct),
-               data.frame(b = c("big", "small"),
+               data.frame(a = c("big", "small"),
                           `1` = c(4, 1),
                           `2` = c(0, 2),
                           `3` = c(2, 0),
@@ -48,7 +48,7 @@ test_that("totals row and col produce correct results when called together", {
   expect_equal(ct %>%
                  add_totals_col %>%
                  add_totals_row(),
-               data.frame(b = c("big", "small", "Total"),
+               data.frame(a = c("big", "small", "Total"),
                           `1` = c(4, 1, 5),
                           `2` = c(0, 2, 2),
                           `3` = c(2, 0, 2),
@@ -67,3 +67,10 @@ test_that("order doesn't matter when totals row and col are called together", {
                  add_totals_col
   )
 })
+
+test_that("both functions work with a single column", {
+  single_col <- data_frame(a = c(as.Date("2016-01-01"), as.Date("2016-02-03")),
+                         b = c(1, 2))
+  expect_error(single_col %>% add_totals_row(), NA) # from http://stackoverflow.com/a/30068233
+  expect_error(single_col %>% add_totals_row(), NA)
+})