Allow environment lookup

mrc-ide · Jul 13, 2023 · 2757d29 · 2757d29
1 parent fa403cc
commit 2757d29
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 9 deletions.
diff --git a/R/query.R b/R/query.R
@@ -35,9 +35,11 @@ outpack_query <- function(expr, name = NULL, scope = NULL, subquery = NULL) {
     expr_parsed <- query_parse_add_scope(expr_parsed, scope)
   }
 
+  lookup <- query_collect_lookup(expr_parsed, subquery_env)
   info <- list(
     single = is_expr_single_value(expr_parsed, subquery_env),
-    parameters = query_parameters(expr_parsed, subquery_env))
+    parameters = lookup$this,
+    environment = lookup$environment)
 
   ret <- list(value = expr_parsed,
               subquery = as.list(subquery_env),
@@ -369,7 +371,7 @@ query_parse_value <- function(expr, context, subquery_env) {
          name = deparse(expr))
   } else if (is_call(expr, ":")) {
     name <- deparse_query(expr[[2]], NULL)
-    valid <- c("parameter", "this")
+    valid <- c("parameter", "this", "environment")
     if (!(name %in% valid)) {
         query_parse_error(sprintf(
           "Invalid lookup '%s'", name), expr, context)
@@ -420,12 +422,15 @@ add_subquery <- function(name, expr, context, subquery_env) {
 }
 
 
-query_parameters <- function(expr_parsed, subquery_env) {
+query_collect_lookup <- function(expr_parsed, subquery_env) {
   env <- new.env(parent = emptyenv())
-  env$seen <- character()
+  collect <- c("this", "environment")
+  for (nm in collect) {
+    env[[nm]] <- character()
+  }
   f <- function(x) {
-    if (is.recursive(x) && x$type == "lookup" && x$name == "this") {
-      env$seen <- c(env$seen, x$query)
+    if (is.recursive(x) && x$type == "lookup" && x$name %in% collect) {
+      env[[x$name]] <- c(env[[x$name]], x$query)
     } else if (x$type == "subquery") {
       sub <- subquery_env[[x$args$name]]
       if (!is.null(sub)) {
@@ -438,5 +443,5 @@ query_parameters <- function(expr_parsed, subquery_env) {
     }
   }
   f(expr_parsed)
-  unique(env$seen)
+  set_names(lapply(collect, function(nm) unique(env[[nm]])), collect)
 }
diff --git a/tests/testthat/test-query.R b/tests/testthat/test-query.R
@@ -213,12 +213,24 @@ test_that("report on parameters used in the query", {
     outpack_query(x)$info$parameters
   }
   expect_equal(f(quote(latest())), character())
-  expect_equal(f(quote(parameter:a < this:a)), "a")
-  expect_equal(f(quote(parameter:a < this:a && this:a > this:b)),
+  expect_equal(f(quote(parameter:x < this:a)), "a")
+  expect_equal(f(quote(parameter:x < this:a && this:a > this:b)),
                c("a", "b"))
 })
 
 
+test_that("report on environment variables used in the query", {
+  f <- function(x) {
+    outpack_query(x)$info$environment
+ }
+  expect_equal(f(quote(latest())), character())
+  expect_equal(f(quote(parameter:x < environment:a)), "a")
+  expect_equal(
+    f(quote(parameter:x < environment:a && environment:a > environment:b)),
+    c("a", "b"))
+})
+
+
 test_that("include parameters from subqueries too", {
   obj <- outpack_query("latest({B})",
                        subquery = list(B = quote(parameter:x < this:y)))