Preserve type of row names in vec_ptype()

Extracted from r-lib#1020
lionel- · Apr 27, 2020 · 09d151b · 09d151b
1 parent 87f4bd8
commit 09d151b
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 6 deletions.
diff --git a/src/type.c b/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);
 

diff --git a/tests/testthat/test-type.R b/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, ]))
+})
diff --git a/tests/testthat/test-type2.R b/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]))