Add context class and register graph at each verb

DESCRIPTION

@@ -19,7 +19,8 @@ Imports: tibble,
     igraph,
     magrittr,
     utils,
-    rlang
+    rlang,
+    R6
 Remotes: hadley/dplyr,
     hadley/rlang
 URL: https://github.com/thomasp85/tidygraph

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)

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()
+}

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)

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)

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))

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)

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)

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)

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)