change size_by handling

- Dont create a node_size column if size_by is NULL or an unsupported column type is specified. Instead set to default sizing (leafCount) and set as an attribute
metrumresearchgroup · Jan 15, 2025 · e2c0909 · e2c0909
1 parent 531833a
commit e2c0909
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git a/R/model-tree.R b/R/model-tree.R
@@ -177,7 +177,7 @@ model_tree.bbi_log_df <- function(
 
   # Format sizing
   tree_data <- size_tree_by(tree_data, size_by = size_by)
-  node_size <- ifelse(is.null(size_by), "leafCount", "node_size")
+  node_size <- attributes(tree_data)$size_by
 
   # Compile attributes into tooltip
   tree_data <- make_tree_tooltip(tree_data, digits = digits, font_size = font_size)
@@ -745,18 +745,19 @@ size_tree_by <- function(tree_data, size_by = NULL){
       # Set node sizes with NA values (including start node) to mean value
       mean_val <- mean(tree_data$node_size, na.rm = TRUE)
       tree_data$node_size[is.na(tree_data$node_size)] <- mean_val
+      attr(tree_data, "size_by") <- "node_size"
     }else{
       col_class <- class(tree_data[[size_by]])
       cli::cli_warn(
         c(
           "Only numeric columns are supported. Column {.val {size_by}} is {.val {col_class}}",
-          "i" = "Setting node size to constant"
+          "i" = "Setting node size to Default"
         )
       )
-      tree_data$node_size <- 1
+      attr(tree_data, "size_by") <- "leafCount"
     }
   }else{
-    tree_data$node_size <- 1
+    attr(tree_data, "size_by") <- "leafCount"
   }
 
   return(tree_data)

diff --git a/tests/testthat/test-model-tree.R b/tests/testthat/test-model-tree.R
@@ -440,14 +440,14 @@ withr::with_options(list(bbr.bbi_exe_path = read_bbi_path()), {
         'Only numeric columns are supported'
       )
       node_sizes <- get_node_attribute(pl_tree$x$data$children, attr = "SizeOfNode")
-      expect_true(dplyr::n_distinct(node_sizes) == 1)
+      expect_true(dplyr::n_distinct(node_sizes) == 2) # leafCount sizing
       # Check character
       expect_warning(
         pl_tree <- model_tree(log_df2, add_summary = FALSE, size_by = "star"),
         'Only numeric columns are supported'
       )
       node_sizes <- get_node_attribute(pl_tree$x$data$children, attr = "SizeOfNode")
-      expect_true(dplyr::n_distinct(node_sizes) == 1)
+      expect_true(dplyr::n_distinct(node_sizes) == 2) # leafCount sizing
     })
 
     it("static plot", {