Preserve type of row names in vec_ptype() (#1050)

Author: lionel-

src/type.c

@@ -1,8 +1,9 @@
 #include "vctrs.h"
-#include "utils.h"
+#include "arg-counter.h"
 #include "ptype-common.h"
 #include "ptype2.h"
-#include "arg-counter.h"
+#include "type-data-frame.h"
+#include "utils.h"
 
 // Initialised at load time
 static SEXP syms_vec_ptype_finalise_dispatch = NULL;
@@ -11,6 +12,7 @@ static SEXP fns_vec_ptype_finalise_dispatch = NULL;
 
 static inline SEXP vec_ptype_slice(SEXP x, SEXP empty);
 static SEXP s3_type(SEXP x, struct vctrs_arg* x_arg);
+static SEXP df_ptype(SEXP x, bool bare);
 
 // [[ register() ]]
 SEXP vctrs_ptype(SEXP x, SEXP x_arg) {
@@ -32,7 +34,7 @@ SEXP vec_ptype(SEXP x, struct vctrs_arg* x_arg) {
   case vctrs_type_character:   return vec_ptype_slice(x, vctrs_shared_empty_chr);
   case vctrs_type_raw:         return vec_ptype_slice(x, vctrs_shared_empty_raw);
   case vctrs_type_list:        return vec_ptype_slice(x, vctrs_shared_empty_list);
-  case vctrs_type_dataframe:   return bare_df_map(x, &col_ptype);
+  case vctrs_type_dataframe:   return df_ptype(x, true);
   case vctrs_type_s3:          return s3_type(x, x_arg);
   case vctrs_type_scalar:      stop_scalar_type(x, x_arg);
   }
@@ -54,10 +56,10 @@ static inline SEXP vec_ptype_slice(SEXP x, SEXP empty) {
 static SEXP s3_type(SEXP x, struct vctrs_arg* x_arg) {
   switch (class_type(x)) {
   case vctrs_class_bare_tibble:
-    return bare_df_map(x, &col_ptype);
+    return df_ptype(x, true);
 
   case vctrs_class_data_frame:
-    return df_map(x, &col_ptype);
+    return df_ptype(x, false);
 
   case vctrs_class_bare_data_frame:
     Rf_errorcall(R_NilValue, "Internal error: Bare data frames should be handled by `vec_ptype()`");
@@ -77,6 +79,24 @@ static SEXP s3_type(SEXP x, struct vctrs_arg* x_arg) {
   return vec_slice(x, R_NilValue);
 }
 
+SEXP df_ptype(SEXP x, bool bare) {
+  SEXP row_nms = PROTECT(df_rownames(x));
+
+  SEXP ptype = R_NilValue;
+  if (bare) {
+    ptype = PROTECT(bare_df_map(x, &col_ptype));
+  } else {
+    ptype = PROTECT(df_map(x, &col_ptype));
+  }
+
+  if (TYPEOF(row_nms) == STRSXP) {
+    Rf_setAttrib(ptype, R_RowNamesSymbol, vctrs_shared_empty_chr);
+  }
+
+  UNPROTECT(2);
+  return ptype;
+}
+
 static SEXP vec_ptype_finalise_unspecified(SEXP x);
 static SEXP vec_ptype_finalise_dispatch(SEXP x);
 

tests/testthat/test-type-dplyr.R

@@ -1,51 +1,53 @@
 
 # `grouped_df` -------------------------------------------------------
 
+bare_mtcars <- unrownames(mtcars)
+
 test_that("grouped-df is proxied and restored", {
-  gdf <- dplyr::group_by(mtcars, cyl)
+  gdf <- dplyr::group_by(bare_mtcars, cyl)
 
   expect_identical(vec_proxy(gdf), gdf)
-  expect_identical(vec_restore(mtcars, gdf), gdf)
+  expect_identical(vec_restore(bare_mtcars, gdf), gdf)
 
   expect_identical(vec_ptype(gdf), gdf[0, ])
 
-  gdf <- dplyr::group_by(mtcars, cyl, am, vs)
+  gdf <- dplyr::group_by(bare_mtcars, cyl, am, vs)
   expect_identical(gdf[0, ], vec_ptype(gdf))
 
-  out <- vec_ptype(dplyr::group_by(mtcars, cyl, .drop = FALSE))
+  out <- vec_ptype(dplyr::group_by(bare_mtcars, cyl, .drop = FALSE))
   expect_drop(out, FALSE)
 })
 
 test_that("can take the common type of grouped tibbles and tibbles", {
-  gdf <- dplyr::group_by(mtcars, cyl)
+  gdf <- dplyr::group_by(bare_mtcars, cyl)
   expect_identical(vec_ptype2(gdf, data.frame()), vec_ptype(gdf))
   expect_identical(vec_ptype2(data.frame(), gdf), vec_ptype(gdf))
   expect_identical(vec_ptype2(gdf, tibble()), vec_ptype(gdf))
   expect_identical(vec_ptype2(tibble(), gdf), vec_ptype(gdf))
 
-  gdf_nodrop <- dplyr::group_by(mtcars, cyl, .drop = FALSE)
+  gdf_nodrop <- dplyr::group_by(bare_mtcars, cyl, .drop = FALSE)
   expect_drop(vec_ptype2(gdf, gdf_nodrop), FALSE)
   expect_drop(vec_ptype2(gdf_nodrop, gdf), FALSE)
-  expect_drop(vec_ptype2(gdf_nodrop, mtcars), FALSE)
-  expect_drop(vec_ptype2(mtcars, gdf_nodrop), FALSE)
+  expect_drop(vec_ptype2(gdf_nodrop, bare_mtcars), FALSE)
+  expect_drop(vec_ptype2(bare_mtcars, gdf_nodrop), FALSE)
 })
 
 test_that("the common type of grouped tibbles includes the union of grouping variables", {
-  gdf1 <- dplyr::group_by(mtcars, cyl)
-  gdf2 <- dplyr::group_by(mtcars, am, vs)
+  gdf1 <- dplyr::group_by(bare_mtcars, cyl)
+  gdf2 <- dplyr::group_by(bare_mtcars, am, vs)
   expect_identical(
     vec_ptype2(gdf1, gdf2),
-    vec_ptype(dplyr::group_by(mtcars, cyl, am, vs))
+    vec_ptype(dplyr::group_by(bare_mtcars, cyl, am, vs))
   )
 })
 
 test_that("can cast to and from `grouped_df`", {
-  gdf <- dplyr::group_by(unrownames(mtcars), cyl)
-  input <- mtcars[10]
-  cast_gdf <- dplyr::group_by(vec_cast(mtcars[10], mtcars), cyl)
+  gdf <- dplyr::group_by(unrownames(bare_mtcars), cyl)
+  input <- bare_mtcars[10]
+  cast_gdf <- dplyr::group_by(vec_cast(bare_mtcars[10], bare_mtcars), cyl)
 
   expect_error(
-    vec_cast(input, dplyr::group_by(mtcars["cyl"], cyl)),
+    vec_cast(input, dplyr::group_by(bare_mtcars["cyl"], cyl)),
     class = "vctrs_error_cast_lossy"
   )
 
@@ -54,15 +56,15 @@ test_that("can cast to and from `grouped_df`", {
     cast_gdf
   )
   expect_identical(
-    vec_cast(gdf, mtcars),
-    unrownames(mtcars)
+    vec_cast(gdf, bare_mtcars),
+    unrownames(bare_mtcars)
   )
 
   expect_identical(
     vec_cast(tibble::as_tibble(input), gdf),
     unrownames(cast_gdf)
   )
-  tib <- tibble::as_tibble(mtcars)
+  tib <- tibble::as_tibble(bare_mtcars)
   expect_identical(
     unrownames(vec_cast(gdf, tib)),
     tib
@@ -71,60 +73,60 @@ test_that("can cast to and from `grouped_df`", {
 
 test_that("casting to `grouped_df` doesn't require grouping variables", {
   expect_identical(
-    vec_cast(mtcars[10], dplyr::group_by(mtcars, cyl)),
-    dplyr::group_by(vec_cast(mtcars[10], mtcars), cyl)
+    vec_cast(bare_mtcars[10], dplyr::group_by(bare_mtcars, cyl)),
+    dplyr::group_by(vec_cast(bare_mtcars[10], bare_mtcars), cyl)
   )
 })
 
 test_that("casting to `grouped_df` handles `drop`", {
-  gdf_nodrop <- dplyr::group_by(mtcars, cyl, .drop = FALSE)
-  expect_identical(vec_cast(mtcars, gdf_nodrop), gdf_nodrop)
+  gdf_nodrop <- dplyr::group_by(bare_mtcars, cyl, .drop = FALSE)
+  expect_identical(vec_cast(bare_mtcars, gdf_nodrop), gdf_nodrop)
 })
 
 test_that("can cbind grouped data frames", {
-  gdf <- dplyr::group_by(mtcars[-10], cyl)
-  df <- unrownames(mtcars)[10]
+  gdf <- dplyr::group_by(bare_mtcars[-10], cyl)
+  df <- unrownames(bare_mtcars)[10]
 
   expect_identical(
     unrownames(vec_cbind(gdf, df)),
-    tibble::as_tibble(mtcars)[c(1:9, 11, 10)]
+    tibble::as_tibble(bare_mtcars)[c(1:9, 11, 10)]
   )
 
-  gdf1 <- dplyr::group_by(mtcars[2], cyl)
-  gdf2 <- dplyr::group_by(mtcars[8:9], vs, am)
+  gdf1 <- dplyr::group_by(bare_mtcars[2], cyl)
+  gdf2 <- dplyr::group_by(bare_mtcars[8:9], vs, am)
   expect_identical(
     unrownames(vec_cbind(gdf1, gdf2)),
-    tibble::as_tibble(mtcars)[c(2, 8, 9)]
+    tibble::as_tibble(bare_mtcars)[c(2, 8, 9)]
   )
 })
 
 
 # `rowwise` ----------------------------------------------------------
 
 test_that("rowwise can be proxied and restored", {
-  rww <- dplyr::rowwise(unrownames(mtcars))
+  rww <- dplyr::rowwise(unrownames(bare_mtcars))
 
   expect_identical(vec_proxy(rww), rww)
-  expect_identical(vec_restore(unrownames(mtcars), rww), rww)
+  expect_identical(vec_restore(unrownames(bare_mtcars), rww), rww)
 
   expect_identical(vec_ptype(rww), rww[0, ])
 })
 
 test_that("can take the common type of rowwise tibbles and tibbles", {
-  rww <- dplyr::rowwise(mtcars)
+  rww <- dplyr::rowwise(bare_mtcars)
   expect_identical(vec_ptype2(rww, data.frame()), vec_ptype(rww))
   expect_identical(vec_ptype2(data.frame(), rww), vec_ptype(rww))
   expect_identical(vec_ptype2(rww, tibble()), vec_ptype(rww))
   expect_identical(vec_ptype2(tibble(), rww), vec_ptype(rww))
 })
 
 test_that("can cast to and from `rowwise_df`", {
-  rww <- unrownames(dplyr::rowwise(mtcars))
-  input <- mtcars[10]
-  cast_rww <- dplyr::rowwise(vec_cast(mtcars[10], mtcars))
+  rww <- unrownames(dplyr::rowwise(bare_mtcars))
+  input <- bare_mtcars[10]
+  cast_rww <- dplyr::rowwise(vec_cast(bare_mtcars[10], bare_mtcars))
 
   expect_error(
-    vec_cast(input, dplyr::rowwise(mtcars["cyl"])),
+    vec_cast(input, dplyr::rowwise(bare_mtcars["cyl"])),
     class = "vctrs_error_cast_lossy"
   )
 
@@ -133,23 +135,23 @@ test_that("can cast to and from `rowwise_df`", {
     cast_rww
   )
   expect_identical(
-    vec_cast(rww, mtcars),
-    unrownames(mtcars)
+    vec_cast(rww, bare_mtcars),
+    unrownames(bare_mtcars)
   )
 
   expect_identical(
     vec_cast(tibble::as_tibble(input), rww),
     unrownames(cast_rww)
   )
-  tib <- tibble::as_tibble(mtcars)
+  tib <- tibble::as_tibble(bare_mtcars)
   expect_identical(
     unrownames(vec_cast(rww, tib)),
     tib
   )
 })
 
 test_that("can cbind rowwise data frames", {
-  df <- unrownames(mtcars)
+  df <- unrownames(bare_mtcars)
   rww <- dplyr::rowwise(df[-2])
   gdf <- dplyr::group_by(df[2], cyl)
 
@@ -162,7 +164,7 @@ test_that("can cbind rowwise data frames", {
 
 test_that("no common type between rowwise and grouped data frames", {
   expect_df_fallback(
-    out <- vec_ptype_common_fallback(dplyr::rowwise(mtcars), dplyr::group_by(mtcars, cyl))
+    out <- vec_ptype_common_fallback(dplyr::rowwise(bare_mtcars), dplyr::group_by(bare_mtcars, cyl))
   )
-  expect_identical(out, tibble::as_tibble(mtcars[0, ]))
+  expect_identical(out, tibble::as_tibble(bare_mtcars[0, ]))
 })

tests/testthat/test-type.R

@@ -221,3 +221,11 @@ test_that("vec_ptype_finalise() requires vector types", {
   expect_error(vec_ptype_finalise(quote(name)), class = "vctrs_error_scalar_type")
   expect_error(vec_ptype_finalise(foobar()), class = "vctrs_error_scalar_type")
 })
+
+# This might change in the future if we decide that prototypes don't
+# have names
+test_that("vec_ptype() preserves type of names and row names", {
+  expect_identical(vec_ptype(c(foo = 1)), named(dbl()))
+  expect_identical(vec_ptype(mtcars), mtcars[0, ])
+  expect_identical(vec_ptype(foobar(mtcars)), foobar(mtcars[0, ]))
+})

tests/testthat/test-type2.R

@@ -296,7 +296,6 @@ test_that("vec_ptype2() methods get prototypes", {
   expect_identical(x, foobar(int()))
   expect_identical(y, foobar(chr()))
 
-  skip("Figure out what to do with row names in `vec_ptype()`")
   vec_ptype2(foobar(mtcars), foobar(iris))
   expect_identical(x, foobar(mtcars[0, , drop = FALSE]))
   expect_identical(y, foobar(iris[0, , drop = FALSE]))