major update for melt_list

rpkgs · Jun 20, 2023 · 37ba706 · 37ba706
1 parent 75cef6e
commit 37ba706
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 17 deletions.
diff --git a/R/melt_list.R b/R/melt_list.R
@@ -20,37 +20,59 @@ melt_list <- function(list, ..., na.rm = TRUE) {
   n <- length(list)
 
   params <- list(...)
-  key <- names(params)[1]
-  vals <- params[[1]]
-  if (is.null(key)) {
-    key <- vals # variable name
-    vals <- names(list)
-  }
-  if (is.null(vals)) vals <- seq_along(list)
-  if (length(vals) == 1) vals <- rep(vals, n)
-  if (is.character(vals)) {
-    if (is_num_char(vals)) {
-      vals %<>% as.numeric()
-    } else {
-      vals %<>% as.factor()
+
+  # check keys and values
+  nkey = length(params)  
+  l_vals = rep(list(NULL), nkey)
+
+  for (k in 1:nkey) {
+    key <- names(params)[k]
+    vals <- params[[k]]
+
+    if (is.null(key)) {
+      key <- vals # variable name
+      vals <- names(list)
     }
+    vals %<>% check_vals(n)
+
+    l_vals[[k]] = vals
+    names(l_vals)[k] = key
   }
 
   first <- list[[1]]
   if (is.data.frame(first)) {
     for (i in seq_along(list)) {
       x <- list[[i]]
-      eval(parse(text = sprintf("x$%s <- vals[i]", key)))
+
+      for (k in 1:nkey) {
+        vals = l_vals[[k]]
+        key = names(l_vals)[k]
+        eval(parse(text = sprintf("x$%s <- vals[i]", key)))
+      }
       list[[i]] <- x
     }
-    # res <- do.call(rbind, list) %>% data.table() # return
     res <- rbindlist(list)
     # } else {
     #     id.vars <- colnames(first)
     #     res <- data.table::melt(list, ..., id.vars = id.vars, na.rm = na.rm)
-    #     colnames(res) <- c(id.vars, key)
+    #     colnames(res) <- c(id.vars, keys)
+  }
+  keys = names(l_vals)
+  res %>% dplyr::relocate(all_of(keys))
+}
+
+# n: the number of variables
+check_vals <- function(vals, n) {
+  if (is.null(vals)) vals <- 1:n
+  if (length(vals) == 1) vals <- rep(vals, n)
+  if (is.character(vals)) {
+    if (is_num_char(vals)) {
+      vals %<>% as.numeric()
+    } else {
+      vals %<>% as.factor()
+    }
   }
-  res %>% dplyr::relocate(key)
+  vals
 }
 
 #' @rdname melt_list

diff --git a/tests/testthat/test-melt_list.R b/tests/testthat/test-melt_list.R
@@ -5,6 +5,8 @@ test_that("melt_list works", {
         df <- data.frame(year = 2010, day = 1:2, month = 1, site = "A")
         l  <- list(a = df, b = df)
         melt_list(l, "id")
+        melt_list(l, x = c(1, 2))
+        melt_list(l, x = c(1, 2), y = c(2, 3))
 
         l2 <- listk("type1" = l, "type2" = l)
         melt_tree(l2, c("type", "id"))