From 8778647b0e3bcd375b33faaf6348deb3a2dc58d6 Mon Sep 17 00:00:00 2001
From: Emi Tanaka <dr.emi.tanaka@gmail.com>
Date: Sat, 25 Nov 2023 14:21:04 +1100
Subject: [PATCH] modularise set_rcrds and expect_rcrds

---
 R/rcrds.R | 23 +++++++++++++++++------
 R/utils.R |  2 +-
 2 files changed, 18 insertions(+), 7 deletions(-)

diff --git a/R/rcrds.R b/R/rcrds.R
index 89c82bd..1cfb8a5 100644
--- a/R/rcrds.R
+++ b/R/rcrds.R
@@ -19,10 +19,10 @@
 #'   set_rcrds_of(unit = "y")
 #' @return An edibble design.
 #' @export
-set_rcrds <- function(.edibble, ...,
+set_rcrds <- function(.edibble = NULL, ...,
                       .name_repair = c("check_unique", "unique", "universal", "minimal"),
                       .record = TRUE) {
-
+  if(is.null(.edibble)) return(structure(match.call(), class = c("edbl_fn", "edbl")))
   not_edibble(.edibble)
   des <- edbl_design(.edibble)
   prov <- activate_provenance(des)
@@ -56,7 +56,8 @@ set_rcrds <- function(.edibble, ...,
       if(!arcrd %in% names(.edibble)) {
         uid <- prov$mapping_to_unit(id = prov$fct_id(name = arcrd))
         uname <- prov$fct_names(id = uid)
-        uids <- prov$lvl_id(value = .edibble[[uname]], fid = uid)
+        uids <- prov$lvl_id(value = attr(.edibble[[uname]], "label-non-nested") %||% .edibble[[uname]],
+                            fid = uid)
         .edibble[[arcrd]] <- new_edibble_rcrd(rep(NA_real_, nrow(.edibble)), uids)
       } else {
         .edibble[[arcrd]] <- new_edibble_rcrd(.edibble[[arcrd]])
@@ -69,7 +70,7 @@ set_rcrds <- function(.edibble, ...,
 
 #' @rdname set_rcrds
 #' @export
-set_rcrds_of <- function(.edibble, ...) {
+set_rcrds_of <- function(.edibble = NULL, ...) {
   unit2rcrd <- list2(...)
   units <- names(unit2rcrd)
   args <- list()
@@ -95,8 +96,18 @@ set_rcrds_of <- function(.edibble, ...) {
 #'   expect_rcrds(y > 0)
 #' @return An edibble design.
 #' @export
-expect_rcrds <- function(.edibble, ..., .record = TRUE) {
-  not_edibble(.edibble)
+expect_rcrds <- function(.edibble = NULL, ..., .record = TRUE) {
+  arg1 <- enquo(.edibble)
+  arg1 <- tryCatch(rlang::eval_tidy(arg1),
+           error = function(e) arg1)
+  if(is.null(arg1)) return(structure(match.call(), class = c("edbl_fn", "edbl")))
+  if(!is_edibble(arg1)) {
+    cl <- match.call()
+    ncl <- length(cl)
+    cl[3:(ncl + 1)] <- cl[2:ncl]
+    cl$.edibble <- NULL
+    return(structure(cl, class = c("edbl_fn", "edbl")))
+  }
   prov <- activate_provenance(.edibble)
   if(.record) prov$record_step()
   dots <- enquos(...)
diff --git a/R/utils.R b/R/utils.R
index 83107d9..53098fa 100644
--- a/R/utils.R
+++ b/R/utils.R
@@ -359,7 +359,7 @@ print.edbl_fct <- function(x, ...) {
     cli::cli_abort(c("Cannot use {.code +} with a single argument",
                      i = "Did you accidentally put {.code +} on a new line?"))
   }
-  if(inherits(e1, "edbl_design")) {
+  if(is_edibble_design(e1) | is_edibble_table(e1)) {
     prov1 <- activate_provenance(e1)
     if(is_edibble_design(e2)) {
       prov2 <- activate_provenance(e2)