From 9ec88f6d8edcae37b1e074a3cf48ab01a7447a91 Mon Sep 17 00:00:00 2001
From: Tomasz Kalinowski <kalinowskit@gmail.com>
Date: Thu, 10 Mar 2022 12:31:12 -0500
Subject: [PATCH] substitute instead of mask

---
 R/zzz-autograph.R | 37 ++++++++++++++++++++++++++++++++++---
 1 file changed, 34 insertions(+), 3 deletions(-)

diff --git a/R/zzz-autograph.R b/R/zzz-autograph.R
index d868cab..de41021 100644
--- a/R/zzz-autograph.R
+++ b/R/zzz-autograph.R
@@ -13,6 +13,16 @@ ag_mask_list <- list(
 )
 
 
+ag_mask_subtitution_list <- alist(
+  `if` = tfautograph::ag_if,
+  `for` = tfautograph::ag_for,
+  `while` = tfautograph::ag_while,
+  `break` = tfautograph::ag_break,
+  `next` = tfautograph::ag_next,
+  `stopifnot` = tfautograph::ag_stopifnot,
+  `on.exit` = tfautograph::ag_on.exit
+)
+
 
 
 #' Autograph R code
@@ -40,16 +50,36 @@ autograph <- function(x) {
   xe <- substitute(x)
   env <- parent.frame()
 
+
   if (is.symbol(xe)) {
-    # function, formula, or something with `environment<-` method
-    environment(x) <- new_ag_mask(parent = environment(x))
+    # function, language, or something with `environment<-` method
+
+    if (inherits(xe, "python.builtin.object"))
+      stop("Only R objects can be autographed")
+
+    if (is.function(x))
+      body(x) <- eval(call("substitute", body(x), ag_mask_subtitution_list))
+    else if (is.language(x)) # formula, expression(), or call() (probably to `{`)
+      x <- eval(call("substitute", x, ag_mask_subtitution_list))
+    else
+      environment(x) <- new_ag_mask(parent = environment(x))
+
+    class(x) <- unique(c("tfautographed", class(x)))
     return(x)
   }
 
   # in line expression
-  fn <- as_outcome_fn(xe, new_ag_mask(parent = env))
+  # fn <- as_outcome_fn(xe, new_ag_mask(parent = env))
+
+  xe <- eval(call("substitute", xe, ag_mask_subtitution_list))
+  fn <- as_outcome_fn(xe, env)
   outcome <- fn()
 
+  outcome <- rapply(list(outcome), function(x) {
+    class(x) <- unique(c("tfautographed", class(x)))
+    x
+  }, how = "replace", classes = c("function", "language"))[[1L]]
+
   export_modified(outcome$modified, env)
 
   if(isFALSE(outcome$visible) ||
@@ -60,6 +90,7 @@ autograph <- function(x) {
 }
 
 
+
 new_ag_mask <- function(parent = parent.frame()) {
 
   ag_mask <- list2env(ag_mask_list, parent = parent)