Search code examples
rr-s3

Function for columns to inherit custom class of data frame in R


I have a list of dataframes

dd <- list()
dd$dat <- list(
  one = data.frame(a = c(1), b = c(2)),
  two = data.frame(c = c(3), d = c(4)),
  three = data.frame(e = c(5), f = c(6))
)

and wrote a function to append a custom class to each dataframe:

# append classes 
append_classes <- function(x, nm) {
  class(x) <- 
    case_when(
      nm == "one" ~ c(class(x), "foo"),
      nm == "two" ~ c(class(x), "bar"),
      TRUE ~ c(class(x), "custom")
    )
  return(x)
}

dd$dat <- imap(dd$dat, append_classes)
class(dd$dat[[1]])

It works!

[1] "data.frame" "foo"  

But now I want to use class inheritance to have the columns in each data frame inherit the foo, bar, and custom classes respectively - how would I do this?

Desired Output

class(dd$dat$one$a)
[1] "numeric" "foo"
class(dd$dat$two$d)
[1] "numeric" "bar"

I'm very new to using S3 so any help is appreciated!!


Solution

  • We can use the imap recursively or use map inside

    library(purrr)
    dd$dat <-  imap(dd$dat, ~ {nm1 <- .y
           map_dfr(append_classes(.x, nm1), ~ append_classes(.x, nm1))
           })
    
    class(dd$dat$one$a)
    #[1] "numeric" "foo"    
    class(dd$dat$two$d)
    #[1] "numeric" "bar" 
    

    Or this can be done with base R using Map/lapply

    dd$dat <- Map(function(x, y) {
         tmp <- append_classes(x, y)
        tmp[] <- lapply(tmp, append_classes, nm = y)
        tmp} , dd$dat, names(dd$dat))
    class(dd$dat$one$a)
    #[1] "numeric" "foo"