printing

NAMESPACE

@@ -2,6 +2,9 @@
 
 S3method("[",causal_tbl)
 S3method("names<-",causal_tbl)
+S3method(ctl_new_pillar,causal_tbl)
+S3method(tbl_format_header,causal_tbl)
+S3method(tbl_format_setup,causal_tbl)
 S3method(tbl_sum,causal_tbl)
 export(as_causal_tbl)
 export(causal_tbl)
@@ -18,5 +21,8 @@ export(set_treatment)
 importFrom(cli,cli_abort)
 importFrom(cli,cli_inform)
 importFrom(cli,cli_warn)
+importFrom(pillar,ctl_new_pillar)
+importFrom(pillar,tbl_format_header)
+importFrom(pillar,tbl_format_setup)
 importFrom(pillar,tbl_sum)
 importFrom(tidyselect,enquo)

R/causal_cols.R

@@ -102,7 +102,7 @@ has_treatment <- function(data) {
 #'   id = c("a", "a", "a", "a", "b", "b", "b", "b"),
 #'   year = rep(2015:2018, 2),
 #'   trt = c(0, 0, 0, 0, 0, 0, 1, 1),
-#'   y = c(1, 3, 2, 3, 1, 3, 4, 5)
+#'   y = c(1, 3, 2, 3, 2, 4, 4, 5)
 #' ) |>
 #'   set_panel(unit=id, time=year)
 #' print(data) # a causal_tbl

R/causal_tbl.R

@@ -104,10 +104,10 @@ reconstruct.causal_tbl <- function(data, old) {
 causal_tbl <- function(..., .outcome=NULL, .treatment=NULL) {
     data = vctrs::df_list(...)
     if (!missing(.outcome)) {
-        outcome <- single_col_name(enquo(.outcome), data, "outcome")
+        .outcome <- single_col_name(enquo(.outcome), data, "outcome")
     }
     if (!missing(.treatment)) {
-        treatment <- single_col_name(enquo(.treatment), data, "treatment")
+        .treatment <- single_col_name(enquo(.treatment), data, "treatment")
     }
 
     new_causal_tbl(data, .outcome=.outcome, .treatment=.treatment)
@@ -179,7 +179,51 @@ assert_df <- function(data, arg) {
 #' @method tbl_sum causal_tbl
 #' @export
 tbl_sum.causal_tbl <- function(x, ...) {
-    lines <- c("A causal_tbl" = pillar::dim_desc(x),
+    lines <- c("A causal_tbl " = pillar::dim_desc(x),
                NextMethod())
     lines[-2] # remove tbl line
 }
+
+#' @importFrom pillar tbl_format_header
+#' @method tbl_format_header causal_tbl
+#' @export
+tbl_format_header.causal_tbl <- function(x, setup, ...) {
+    default_header <- NextMethod()
+    new_header <- paste0(
+        pillar::style_subtle(cli::format_inline("# A {.cls causal_tbl}")),
+        pillar::style_subtle(paste0(" [", pillar::dim_desc(x), "]"))
+    )
+    c(new_header, default_header[-1])
+}
+
+#' @importFrom pillar tbl_format_setup
+#' @method tbl_format_setup causal_tbl
+#' @export
+tbl_format_setup.causal_tbl <- function(x, width, ..., n, max_extra_cols, max_footer_lines, focus) {
+    NextMethod(focus=unlist(causal_cols(x)))
+}
+
+
+#' @importFrom pillar ctl_new_pillar
+#' @method ctl_new_pillar causal_tbl
+#' @export
+ctl_new_pillar.causal_tbl <- function(controller, x, width, ..., title = NULL) {
+    out <- NextMethod()
+    cols <- causal_cols(controller)
+    marker_type = names(cols)[match(title, cols)]
+    marker = if (length(marker_type) == 0 || is.na(marker_type)) {
+        ""
+    } else {
+        pillar::style_subtle(c(
+            outcome="[out]", treatment="[trt]",
+            panel_unit="[unit]", panel_time="[time]"
+        )[marker_type])
+    }
+
+    pillar::new_pillar(list(
+        marker = pillar::new_pillar_component(list(marker), width = nchar(marker)),
+        title = out$title,
+        type = out$type,
+        data = out$data
+    ))
+}

R/dplyr.R

@@ -1,11 +1,11 @@
 group_by.causal_tbl <- function(.data, ..., .add=FALSE) {
-    reconstruct.causal_tbl(NextMethod(), data)
+    reconstruct.causal_tbl(NextMethod(), .data)
 }
 ungroup.causal_tbl <- function(.data, ...) {
-    reconstruct.causal_tbl(NextMethod(), data)
+    reconstruct.causal_tbl(NextMethod(), .data)
 }
 rowwise.causal_tbl <- function(.data, ...) {
-    reconstruct.causal_tbl(NextMethod(), data)
+    reconstruct.causal_tbl(NextMethod(), .data)
 }
 
 dplyr_reconstruct.causal_tbl <- function(data, template) {

tests/testthat/test-causal_cols.R

@@ -49,14 +49,19 @@ test_that("getting and setting panel data", {
         id = c("a", "a", "a", "a", "b", "b", "b", "b"),
         year = rep(2015:2018, 2),
         trt = c(0, 0, 0, 0, 0, 0, 1, 1),
-        y = c(1, 3, 2, 3, 1, 3, 4, 5)
+        y = c(1, 3, 2, 3, 2, 4, 4, 5)
     )
 
     x_panel = set_panel(x, unit=id, time=year)
     expect_equal(get_panel(x_panel), list(unit = "id", time = "year"))
     expect_true(has_panel(x_panel))
     expect_no_error(validate_causal_tbl(x_panel))
 
+    x_panel2 = set_treatment(set_outcome(x_panel, y), trt)
+    expect_true(has_panel(x_panel2))
+    expect_true(has_treatment(x_panel2))
+    expect_true(has_outcome(x_panel2))
+
     expect_error(set_panel(x, unit=id), "for `time`")
     expect_error(set_panel(x, time=year), "for `unit`")
     expect_error(set_panel(x, unit=year, time=id), "convert `id`")