From c26426a4862b6786175aa58fc69376116f2c7033 Mon Sep 17 00:00:00 2001
From: Lionel Henry <lionel.hry@gmail.com>
Date: Wed, 27 Feb 2019 10:21:34 +0100
Subject: [PATCH] Fix `compose()` with generic functions

Closes #629
Closes #639
---
 NEWS.md                       |  3 ++
 R/compose.R                   | 12 +++---
 tests/testthat/test-compose.R | 81 +++++++++++++++++++++++++++++++++++
 3 files changed, 91 insertions(+), 5 deletions(-)

diff --git a/NEWS.md b/NEWS.md
index f4ddeb6a..ec445bc7 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -1,6 +1,9 @@
 
 # purrr 0.3.0.9000
 
+* `compose()` now works with generic functions again (#629, #639). Its
+  set of unit tests was expanded to cover many edge cases.
+
 
 # purrr 0.3.0
 
diff --git a/R/compose.R b/R/compose.R
index eeec0124..ea55a937 100644
--- a/R/compose.R
+++ b/R/compose.R
@@ -48,19 +48,21 @@ compose <- function(..., .dir = c("backward", "forward")) {
     fns <- fns[-1]
   }
 
-  body <- expr({
-    out <- !!fn_body(first_fn)
+  composed <- function() {
+    call <- sys.call()
+    call[[1]] <- first_fn
+    out <- eval_bare(call, caller_env())
 
-    fns <- !!fns
     for (fn in fns) {
       out <- fn(out)
     }
 
     out
-  })
+  }
+  formals(composed) <- formals(first_fn)
 
   structure(
-    new_function(formals(first_fn), body, fn_env(first_fn)),
+    composed,
     class = c("purrr_function_compose", "function"),
     first_fn = first_fn,
     fns = fns
diff --git a/tests/testthat/test-compose.R b/tests/testthat/test-compose.R
index b612031a..624e91d6 100644
--- a/tests/testthat/test-compose.R
+++ b/tests/testthat/test-compose.R
@@ -65,3 +65,84 @@ test_that("compose() with 0 inputs returns the identity", {
 test_that("compose() with 1 input is a noop", {
   expect_identical(compose(toupper)(letters), toupper(letters))
 })
+
+test_that("compose() works with generic functions (#629)", {
+  purrr__gen <- function(x) UseMethod("purrr__gen")
+
+  # Can pass lexical context of methods through lambdas
+  local({
+    purrr__gen.default <- function(x) x + 1
+    expect_identical(compose(~ purrr__gen(.x))(0), 1)
+    expect_identical(compose(~ purrr__gen(.x), ~ purrr__gen(.x))(0), 2)
+  })
+
+  # Why doesn't this work locally?
+  scoped_bindings(.env = global_env(),
+    purrr__gen.default = function(x) x + 1
+  )
+
+  expect_identical(compose(purrr__gen)(0), 1)
+  expect_identical(compose(purrr__gen, purrr__gen)(0), 2)
+})
+
+test_that("compose() works with generic functions (#639)", {
+  n_unique <- purrr::compose(length, unique)
+  expect_identical(n_unique(iris$Species), 3L)
+})
+
+test_that("compose() works with argument matching functions", {
+  # They inspect their dynamic context via sys.function()
+  fn <- function(x = c("foo", "bar")) match.arg(x)
+  expect_identical(compose(fn)("f"), "foo")
+  expect_identical(compose(fn, fn)("f"), "foo")
+})
+
+test_that("compose() works with non-local exits", {
+  fn <- function(x) return(x)
+  expect_identical(compose(fn)("foo"), "foo")
+  expect_identical(compose(fn, fn)("foo"), "foo")
+  expect_identical(compose(~ return(paste(.x, "foo")), ~ return("bar"))(), "bar foo")
+})
+
+test_that("compose() preserves lexical environment", {
+  fn <- local({
+    `_foo` <- "foo"
+    function(...) `_foo`
+  })
+  expect_identical(compose(fn)(), "foo")
+  expect_identical(compose(fn, fn)(), "foo")
+})
+
+test_that("compose() can take dots from multiple environments", {
+  f <- function(...) {
+    `_foo` <- "foo"
+    g(`_foo`, ...)
+  }
+  g <- function(...) {
+    `_bar` <- "bar"
+    h(`_bar`, ...)
+  }
+  h <- function(...) {
+    `_baz` <- "baz"
+    fn(`_baz`, ...)
+  }
+  `_quux` <- "quux"
+
+  # By value
+  fn <- compose(function(...) c(...))
+  expect_identical(f(`_quux`), c("baz", "bar", "foo", "quux"))
+
+  # By expression (base)
+  fn <- compose(function(...) substitute(...()))
+  expect_identical(f(`_quux`), as.pairlist(exprs(`_baz`, `_bar`, `_foo`, `_quux`)))
+
+  # By expression (rlang)
+  fn <- compose(function(...) enquos(...))
+  quos <- f(`_quux`)
+
+  frame <- current_env()
+  expect_true(is_reference(quo_get_env(quos[[4]]), frame))
+  expect_false(is_reference(quo_get_env(quos[[3]]), frame))
+
+  expect_identical(unname(map_chr(quos, as_name)), c("_baz", "_bar", "_foo", "_quux"))
+})