Better subsetting for tbls.

Closes tidyverse#587. Closes tidyverse#610

NAMESPACE

@@ -1,6 +1,7 @@
 # Generated by roxygen2 (4.0.2): do not edit by hand
 
 S3method("[",grouped_df)
+S3method("[",tbl_df)
 S3method(all.equal,tbl_df)
 S3method(all.equal,tbl_dt)
 S3method(anti_join,data.frame)

R/tbl-df.r

@@ -5,6 +5,16 @@
 #' tbl objects only print a few rows and all the columns that fit on one
 #' screen, describing the rest of it as text.
 #'
+#' @section Methods:
+#'
+#' \code{tbl_df} implements two important base methods:
+#'
+#' \describe{
+#' \item{print}{Only prints the first 10 rows, and the columns that fit on
+#'   screen}
+#' \item{\code{[}}{Never simplifies (drops), so always returns data.frame}
+#' }
+#'
 #' @export
 #' @param data a data frame
 #' @examples
@@ -93,3 +103,32 @@ print.tbl_df <- function(x, ..., n = NULL) {
   trunc_mat(x, n = n)
 }
 
+#' @export
+`[.tbl_df` <- function (x, i, j, drop = FALSE) {
+  if (missing(i) && missing(j)) return(x)
+  if (drop) warning("drop ignored", call. = FALSE)
+
+  # Escape early if nargs() == 2L; ie, column subsetting
+  if (nargs() == 2L) {
+    result <- .subset(x, i)
+    class(result) <- c("tbl_df", "data.frame")
+    attr(result, "row.names") <- .set_row_names(length(x[[1]]))
+    return(result)
+  }
+
+  # First, subset columns
+  if (!missing(j)) {
+    x <- .subset(x, j)
+  }
+
+  # Next, subset rows
+  if (!missing(i)) {
+    x <- lapply(x, `[`, i)
+  }
+
+  # TODO: handle 0 column case
+  class(x) <- c("tbl_df", "data.frame")
+  attr(x, "row.names") <- .set_row_names(length(x[[1]]))
+  x
+}
+

man/tbl_df.Rd

@@ -14,6 +14,17 @@ a \code{tbl_df} over a regular data frame is the printing:
 tbl objects only print a few rows and all the columns that fit on one
 screen, describing the rest of it as text.
 }
+\section{Methods}{
+
+
+\code{tbl_df} implements two important base methods:
+
+\describe{
+\item{print}{Only prints the first 10 rows, and the columns that fit on
+  screen}
+\item{\code{[}}{Never simplifies (drops), so always returns data.frame}
+}
+}
 \examples{
 ds <- tbl_df(mtcars)
 ds

tests/testthat/test-tbl-df.r

@@ -3,4 +3,11 @@ context("tbl_df/grouped_df")
 test_that("grouped_df returns tbl_df if no groups", {
   df <- grouped_df(mtcars, list())
   expect_equal(class(df), c("tbl_df", "tbl", "data.frame"))
-})
+})
+
+test_that("[ never drops", {
+  mtcars2 <- tbl_df(mtcars)
+  expect_is(mtcars2[, 1], "data.frame")
+  expect_is(mtcars2[, 1], "tbl_df")
+  expect_equal(mtcars2[, 1], mtcars2[1])
+})