Fix col_labels bug for named label attributes (#301)

# Pull Request  Fixes #300 --------- Signed-off-by: Emily de la Rua <[email protected]> Signed-off-by: Pawel Rucki <[email protected]> Co-authored-by: Pawel Rucki <[email protected]>
insightsengineering · Mar 1, 2024 · 6da69db · 6da69db
1 parent dd959ad
commit 6da69db
Show file tree

Hide file tree

Showing 4 changed files with 113 additions and 9 deletions.
diff --git a/NEWS.md b/NEWS.md
@@ -1,5 +1,11 @@
 # teal.data 0.5.0.9001
 
+### Enhancements
+* `col_relabel` supports `NA` to remove labels (similar to the `col_labels<-`).
+
+### Bug fixes
+* Fixed bug in `col_labels` causing incorrect label names to be returned when input data contains named label attributes.
+
 # teal.data 0.5.0
 
 ### Bug fixes

diff --git a/R/formatters_var_labels.R b/R/formatters_var_labels.R
@@ -11,7 +11,7 @@
 #'  if named, names must match variable names in `x` and will be used as key to set labels;
 #'  use `NA` to remove label from variable
 #' @param ... name-value pairs, where name corresponds to a variable name in `x`
-#'  and value is the new variable label
+#'  and value is the new variable label; use `NA` to remove label from variable
 #'
 #' @return
 #' For `col_labels`, named character vector of variable labels, the names being the corresponding variable names.
@@ -43,7 +43,18 @@ col_labels <- function(x, fill = FALSE) {
     return(character(0L))
   }
 
-  labels <- lapply(x, attr, "label")
+  labels <- sapply(x, function(i) as.vector(attr(i, "label", exact = TRUE)), simplify = FALSE, USE.NAMES = TRUE)
+  mapply(
+    function(name, label) {
+      checkmate::assert_string(
+        label,
+        .var.name = sprintf("\"label\" attribute of column \"%s\"", name),
+        null.ok = TRUE
+      )
+    },
+    name = names(x),
+    label = labels
+  )
 
   nulls <- vapply(labels, is.null, logical(1L))
   if (any(nulls)) {
@@ -71,12 +82,24 @@ col_labels <- function(x, fill = FALSE) {
   varnames <-
     if (is.null(names(value))) {
       names(x)
+    } else if (any(names(value) == "")) {
+      specified_cols <- names(value)[names(value) != ""]
+      checkmate::assert_subset(specified_cols, names(x), .var.name = "names of value")
+      res <- names(value)
+      res[res == ""] <- setdiff(names(x), specified_cols)
+      res
     } else {
-      checkmate::assert_set_equal(names(value), names(x), .var.name = "column names")
+      checkmate::assert_set_equal(names(value), names(x), .var.name = "names of value")
       names(value)
     }
 
-  x[varnames] <- mapply(`attr<-`, x = x[varnames], which = "label", value = value, SIMPLIFY = FALSE)
+  for (i in seq_along(value)) {
+    if (is.na(value[i])) {
+      attr(x[[varnames[i]]], "label") <- NULL
+    } else {
+      attr(x[[varnames[i]]], "label") <- value[[i]]
+    }
+  }
   x
 }
 
@@ -91,8 +114,14 @@ col_relabel <- function(x, ...) {
   varnames <- names(value)
 
   checkmate::assert_subset(varnames, names(x), .var.name = "names of ...")
-  lapply(value, checkmate::assert_string, .var.name = "element of ...")
+  lapply(value, checkmate::assert_string, .var.name = "element of ...", na.ok = TRUE)
 
-  x[varnames] <- mapply(`attr<-`, x = x[varnames], which = "label", value = value, SIMPLIFY = FALSE)
+  for (i in seq_along(value)) {
+    if (is.na(value[i])) {
+      attr(x[[varnames[i]]], "label") <- NULL
+    } else {
+      attr(x[[varnames[i]]], "label") <- value[[i]]
+    }
+  }
   x
 }
diff --git a/man/col_labels.Rd b/man/col_labels.Rd
diff --git a/tests/testthat/test-formatters_var_labels.R b/tests/testthat/test-formatters_var_labels.R
@@ -21,8 +21,39 @@ testthat::test_that("col_labels returns a vector of column names when fill = TRU
   )
 })
 
+testthat::test_that("col_labels works with labels having additional attributes (including names)", {
+  iris_df <- utils::head(iris, 2)
+  attr(iris_df$Species, "label") <- structure("Label for Species", names = "blah", foo = "bar")
+  testthat::expect_identical(col_labels(iris_df)[["Species"]], "Label for Species")
+})
 
-# col_labels ----
+testthat::test_that("col_labels returns only 'names' attribute and ignores all the rest", {
+  iris_df <- utils::head(iris, 2)
+  attr(iris_df$Species, "label") <- structure("Label for Species", names = "blah", foo = "bar")
+  testthat::expect_identical(attributes(col_labels(iris_df)["Species"]), list(names = "Species"))
+})
+
+testthat::test_that("col_labels throws if label is not a character", {
+  iris_df <- utils::head(iris, 2)
+  attr(iris_df$Species, "label") <- structure(1, names = "blah", foo = "bar")
+  testthat::expect_error(
+    col_labels(iris_df),
+    "Assertion on '\"label\" attribute of column \"Species\"' failed",
+    fixed = TRUE
+  )
+})
+
+testthat::test_that("col_labels throws if label is not a character of length 1", {
+  iris_df <- utils::head(iris, 2)
+  attr(iris_df$Species, "label") <- structure(c("a", "b"), names = "blah", foo = "bar")
+  testthat::expect_error(
+    col_labels(iris_df),
+    "Assertion on '\"label\" attribute of column \"Species\"' failed",
+    fixed = TRUE
+  )
+})
+
+# col_labels<- ----
 testthat::test_that("col_labels<- value accepts character vector", {
   iris_df <- utils::head(iris, 2)
   testthat::expect_error(
@@ -49,7 +80,7 @@ testthat::test_that("col_labels<- value names must be same as variable names", {
   iris_df <- utils::head(iris, 2)
   testthat::expect_error(
     col_labels(iris_df) <- stats::setNames(as.character(1:5), toupper(names(iris_df))),
-    "Assertion on 'column names' failed: Must be a permutation of set"
+    "Assertion on 'names of value' failed: Must be a permutation of set"
   )
 })
 
@@ -63,6 +94,24 @@ testthat::test_that("col_labels<- sets variable labels when passed unnamed chara
   )
 })
 
+testthat::test_that("col_labels<- sets variable labels when passed partially named character vector", {
+  x <- data.frame(a = 1, b = 2, c = 3)
+  col_labels(x) <- c(a = "A", "B", "C")
+  testthat::expect_identical(
+    col_labels(x),
+    c(a = "A", b = "B", c = "C")
+  )
+})
+
+testthat::test_that("col_labels<- sets variable labels when passed partially named, unordered character vector", {
+  x <- data.frame(a = 1, b = 2, c = 3)
+  col_labels(x) <- c(b = "B", "A", "C")
+  testthat::expect_identical(
+    col_labels(x),
+    c(a = "A", b = "B", c = "C")
+  )
+})
+
 testthat::test_that("col_labels<- sets variable labels when passed named character vector", {
   iris_df <- utils::head(iris, 2)
   labels <- stats::setNames(paste("label for", names(iris_df)), names(iris_df))
@@ -83,6 +132,16 @@ testthat::test_that("col_labels<- matches labels to variables by names of values
   )
 })
 
+testthat::test_that("col_labels<- removes labels on NA_character_", {
+  x <- data.frame(a = 1, b = 2, c = 3)
+  col_labels(x) <- c("A", "B", "C")
+  col_labels(x) <- c(b = NA, "AA", NA)
+  testthat::expect_identical(
+    col_labels(x),
+    c(a = "AA", b = NA, c = NA)
+  )
+})
+
 
 # col_relabel ----
 testthat::test_that("col_relabel correctly changes column labels in a data frame", {
@@ -105,3 +164,13 @@ testthat::test_that("col_relabel returns the original data.frame when no new lab
   iris_df <- col_relabel(iris)
   testthat::expect_equal(iris_df, iris)
 })
+
+testthat::test_that("col_relabel removes labels on NA_character_", {
+  x <- data.frame(a = 1, b = 2, c = 3)
+  col_labels(x) <- c("A", "B", "C")
+  x <- col_relabel(x, b = NA_character_)
+  testthat::expect_identical(
+    col_labels(x),
+    c(a = "A", b = NA, c = "C")
+  )
+})