From c8122baee7ffaf51a256627e6e826ecf6ce68e37 Mon Sep 17 00:00:00 2001
From: Thomas Lin Pedersen <thomasp85@gmail.com>
Date: Mon, 13 Mar 2017 22:35:14 +0100
Subject: [PATCH] Add context class and register graph at each verb

---
 DESCRIPTION  |  3 ++-
 R/arrange.R  |  2 ++
 R/context.R  | 49 +++++++++++++++++++++++++++++++++++++++++++++++++
 R/distinct.R |  2 ++
 R/filter.R   |  2 ++
 R/group_by.R |  2 ++
 R/mutate.R   |  2 ++
 R/rename.R   |  2 ++
 R/select.R   |  2 ++
 R/slice.R    |  2 ++
 10 files changed, 67 insertions(+), 1 deletion(-)
 create mode 100644 R/context.R

diff --git a/DESCRIPTION b/DESCRIPTION
index 5c16869..956ad4a 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -19,7 +19,8 @@ Imports: tibble,
     igraph,
     magrittr,
     utils,
-    rlang
+    rlang,
+    R6
 Remotes: hadley/dplyr,
     hadley/rlang
 URL: https://github.com/thomasp85/tidygraph
diff --git a/R/arrange.R b/R/arrange.R
index 8ab3174..533ac44 100644
--- a/R/arrange.R
+++ b/R/arrange.R
@@ -1,6 +1,8 @@
 #' @export
 #' @importFrom dplyr arrange
 arrange.tbl_graph <- function(.data, ...) {
+  .graph_context$set(.data)
+  on.exit(.graph_context$clear())
   d_tmp <- as_tibble(.data)
   if ('.tbl_graph_index' %in% names(d_tmp)) {
     stop('The attribute name ".tbl_graph_index" is reserved', call. = FALSE)
diff --git a/R/context.R b/R/context.R
new file mode 100644
index 0000000..bd18b26
--- /dev/null
+++ b/R/context.R
@@ -0,0 +1,49 @@
+#' @importFrom R6 R6Class
+ContextBuilder <- R6Class(
+  'ContextBuilder',
+  public = list(
+    set = function(graph) {
+      stopifnot(inherits(graph, 'tbl_graph'))
+      private$context <- graph
+      invisible(self)
+    },
+    clear = function() {
+      private$context <- NULL
+    },
+    alive = function() {
+      !is.null(private$context)
+    },
+    graph = function() {
+      private$check()
+      private$context
+    },
+    nodes = function() {
+      as_tibble(self$graph, active = 'nodes')
+    },
+    edges = function() {
+      as_tibble(self$graph, active = 'edges')
+    }
+  ),
+  private = list(
+    context = NULL,
+    check = function() {
+      if (!self$alive()) {
+        stop('This function should not be called directly', call. = FALSE)
+      }
+    }
+  )
+)
+.graph_context <- ContextBuilder$new()
+
+#' @export
+.G <- function() {
+  .graph_context$graph()
+}
+#' @export
+.N <- function() {
+  .graph_context$nodes()
+}
+#' @export
+.E <- function() {
+  .graph_context$edges()
+}
diff --git a/R/distinct.R b/R/distinct.R
index 823fea1..d3e747c 100644
--- a/R/distinct.R
+++ b/R/distinct.R
@@ -3,6 +3,8 @@
 #' @importFrom utils head
 #' @export
 distinct.tbl_graph <- function(.data, ..., .keep_all = FALSE) {
+  .graph_context$set(.data)
+  on.exit(.graph_context$clear())
   d_tmp <- as_tibble(.data)
   if ('.tbl_graph_index' %in% names(d_tmp)) {
     stop('The attribute name ".tbl_graph_index" is reserved', call. = FALSE)
diff --git a/R/filter.R b/R/filter.R
index 945b08f..21a6800 100644
--- a/R/filter.R
+++ b/R/filter.R
@@ -2,6 +2,8 @@
 #' @importFrom dplyr filter
 #' @importFrom igraph delete_vertices delete_edges
 filter.tbl_graph <- function(.data, ...) {
+  .graph_context$set(.data)
+  on.exit(.graph_context$clear())
   d_tmp <- as_tibble(.data)
   if ('.tbl_graph_index' %in% names(d_tmp)) {
     stop('The attribute name ".tbl_graph_index" is reserved', call. = FALSE)
diff --git a/R/group_by.R b/R/group_by.R
index 2a6f6fe..9f81b5a 100644
--- a/R/group_by.R
+++ b/R/group_by.R
@@ -1,6 +1,8 @@
 #' @importFrom dplyr group_by
 #' @export
 group_by.tbl_graph <- function(.data, ..., add = FALSE) {
+  .graph_context$set(.data)
+  on.exit(.graph_context$clear())
   d_tmp <- as_tibble(.data)
   d_tmp <- group_by(d_tmp, ..., add = add)
   apply_groups(.data, attributes(d_tmp))
diff --git a/R/mutate.R b/R/mutate.R
index 4a21d77..2fecda1 100644
--- a/R/mutate.R
+++ b/R/mutate.R
@@ -1,6 +1,8 @@
 #' @export
 #' @importFrom dplyr mutate
 mutate.tbl_graph <- function(.data, ...) {
+  .graph_context$set(.data)
+  on.exit(.graph_context$clear())
   d_tmp <- as_tibble(.data)
   d_tmp <- mutate(d_tmp, ...)
   set_graph_data(.data, d_tmp)
diff --git a/R/rename.R b/R/rename.R
index 979fb66..0e2fc25 100644
--- a/R/rename.R
+++ b/R/rename.R
@@ -1,6 +1,8 @@
 #' @export
 #' @importFrom dplyr rename
 rename.tbl_graph <- function(.data, ...) {
+  .graph_context$set(.data)
+  on.exit(.graph_context$clear())
   d_tmp <- as_tibble(.data)
   d_tmp <- rename(d_tmp, ...)
   set_graph_data(.data, d_tmp)
diff --git a/R/select.R b/R/select.R
index 95c4b49..a8d49d3 100644
--- a/R/select.R
+++ b/R/select.R
@@ -1,6 +1,8 @@
 #' @export
 #' @importFrom dplyr select
 select.tbl_graph <- function(.data, ...) {
+  .graph_context$set(.data)
+  on.exit(.graph_context$clear())
   d_tmp <- as_tibble(.data)
   d_tmp <- select(d_tmp, ...)
   set_graph_data(.data, d_tmp)
diff --git a/R/slice.R b/R/slice.R
index fbb0120..12e9ddb 100644
--- a/R/slice.R
+++ b/R/slice.R
@@ -2,6 +2,8 @@
 #' @importFrom dplyr slice
 #' @importFrom igraph delete_vertices delete_edges
 slice.tbl_graph <- function(.data, ...) {
+  .graph_context$set(.data)
+  on.exit(.graph_context$clear())
   d_tmp <- as_tibble(.data)
   if ('.tbl_graph_index' %in% names(d_tmp)) {
     stop('The attribute name ".tbl_graph_index" is reserved', call. = FALSE)