setDT() works on S4 slots

Author: MichaelChirico

R/data.table.R

@@ -1221,7 +1221,7 @@ replace_dot_alias = function(e) {
             setalloccol(x, n, verbose=verbose)   # always assigns to calling scope; i.e. this scope
             if (is.name(name)) {
               assign(as.character(name),x,parent.frame(),inherits=TRUE)
-            } else if (name %iscall% c('$', '[[') && is.name(name[[2L]])) {
+            } else if (.is_simple_extraction(name)) { # TODO(#6702): use a helper here as the code is very similar to setDT().
               k = eval(name[[2L]], parent.frame(), parent.frame())
               if (is.list(k)) {
                 origj = j = if (name[[1L]] == "$") as.character(name[[3L]]) else eval(name[[3L]], parent.frame(), parent.frame())
@@ -1233,6 +1233,8 @@ replace_dot_alias = function(e) {
                 .Call(Csetlistelt,k,as.integer(j), x)
               } else if (is.environment(k) && exists(as.character(name[[3L]]), k)) {
                 assign(as.character(name[[3L]]), x, k, inherits=FALSE)
+              } else if (isS4(k)) {
+                .Call(CsetS4elt, k, as.character(name[[3L]]), x)
               }
             } # TO DO: else if env$<- or list$<-
           }
@@ -2967,7 +2969,7 @@ setDT = function(x, keep.rownames=FALSE, key=NULL, check.names=FALSE) {
   if (is.name(name)) {
     name = as.character(name)
     assign(name, x, parent.frame(), inherits=TRUE)
-  } else if (name %iscall% c('$', '[[') && is.name(name[[2L]])) {
+  } else if (.is_simple_extraction(name)) {
     # common case is call from 'lapply()'
     k = eval(name[[2L]], parent.frame(), parent.frame())
     if (is.list(k)) {
@@ -2979,9 +2981,11 @@ setDT = function(x, keep.rownames=FALSE, key=NULL, check.names=FALSE) {
             stopf("Item '%s' not found in names of input list", origj)
         }
       }
-      .Call(Csetlistelt,k,as.integer(j), x)
+      .Call(Csetlistelt, k, as.integer(j), x)
     } else if (is.environment(k) && exists(as.character(name[[3L]]), k)) {
       assign(as.character(name[[3L]]), x, k, inherits=FALSE)
+    } else if (isS4(k)) {
+      .Call(CsetS4elt, k, as.character(name[[3L]]), x)
     }
   }
   .Call(CexpandAltRep, x)  # issue#2866 and PR#2882
@@ -3048,6 +3052,9 @@ rleidv = function(x, cols=seq_along(x), prefix=NULL) {
   is.name(e[[1L]]) && is.name(e[[2L]])         # e.g. V1:V2, but not min(V1):max(V2) or 1:max(V2)
 }
 
+# for assignments like x[[1]][, a := 2] or setDT(x@DT)
+.is_simple_extraction = function(e) e %iscall% c('$', '@', '[[') && is.name(e[[2L]])
+
 # GForce functions
 #   to add a new function to GForce (from the R side -- the easy part!):
 #     (1) add it to gfuns

inst/tests/S4.Rraw

@@ -109,3 +109,11 @@ DT = data.table(a = rep(1:2, c(1, 100)))
 # Set the S4 bit on a simple object
 DT[, b := asS4(seq_len(.N))]
 test(6, DT[, b, by=a, verbose=TRUE][, isS4(b)], output="dogroups: growing")
+
+# setDT() works for a data.frame slot, #6701
+setClass("DataFrame", slots=c(x="data.frame"))
+DF = new("DataFrame", x=data.frame(a=1))
+setDT(DF@x)
+test(7.1, is.data.table(DF@x))
+setClass("DataTable", slots=c(x="data.table"))
+test(7.2, options=c(datatable.alloccol=0L), new("DataTable", x=data.table(a=1))@x[, b := 2L]$b, 2L)

src/data.table.h

@@ -299,6 +299,7 @@ SEXP freadR(SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SE
 SEXP fwriteR(SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP);
 SEXP rbindlist(SEXP, SEXP, SEXP, SEXP, SEXP);
 SEXP setlistelt(SEXP, SEXP, SEXP);
+SEXP setS4elt(SEXP, SEXP, SEXP);
 SEXP address(SEXP);
 SEXP expandAltRep(SEXP);
 SEXP fmelt(SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP, SEXP);

src/init.c

@@ -72,6 +72,7 @@ R_CallMethodDef callMethods[] = {
 {"Crbindlist", (DL_FUNC) &rbindlist, -1},
 {"Cvecseq", (DL_FUNC) &vecseq, -1},
 {"Csetlistelt", (DL_FUNC) &setlistelt, -1},
+{"CsetS4elt", (DL_FUNC) &setS4elt, -1},
 {"Caddress", (DL_FUNC) &address, -1},
 {"CexpandAltRep", (DL_FUNC) &expandAltRep, -1},
 {"Cfmelt", (DL_FUNC) &fmelt, -1},

src/wrappers.c

@@ -59,18 +59,26 @@ SEXP copy(SEXP x)
   return(duplicate(x));
 }
 
+// Internal use only. So that := can update elements of a list of data.table, #2204. Just needed to overallocate/grow the VECSXP.
 SEXP setlistelt(SEXP l, SEXP i, SEXP value)
 {
-  R_len_t i2;
-  // Internal use only. So that := can update elements of a list of data.table, #2204. Just needed to overallocate/grow the VECSXP.
-  if (!isNewList(l)) error(_("First argument to setlistelt must be a list()"));
-  if (!isInteger(i) || LENGTH(i)!=1) error(_("Second argument to setlistelt must a length 1 integer vector"));
-  i2 = INTEGER(i)[0];
+  if (!isNewList(l)) internal_error(__func__, "First argument to setlistelt must be a list()");
+  if (!isInteger(i) || LENGTH(i)!=1) internal_error(__func__, "Second argument to setlistelt must a length 1 integer vector");
+  R_len_t i2 = INTEGER(i)[0];
   if (LENGTH(l) < i2 || i2<1) error(_("i (%d) is outside the range of items [1,%d]"),i2,LENGTH(l));
   SET_VECTOR_ELT(l, i2-1, value);
   return(R_NilValue);
 }
 
+// Internal use only. So that := can update elements of a slot of data.table, #6701.
+SEXP setS4elt(SEXP obj, SEXP name, SEXP value)
+{
+  if (!isS4(obj)) internal_error(__func__, "First argument to setS4elt must be an S4 object");
+  if (!isString(name) || LENGTH(name)!=1) internal_error(__func__, "Second argument to setS4elt must be a character string");
+  R_do_slot_assign(obj, name, value);
+  return(R_NilValue);
+}
+
 SEXP address(SEXP x)
 {
   // A better way than : http://stackoverflow.com/a/10913296/403310