Reputation: 710
I'm writing a short R package which contains a function. The function returns a list of vectors. I would like to use the plot function in order to plot by default a plot done with some of those vectors, add lines and add a new parameter.
As an example, if I use the survival
package I can get the following:
library(survival)
data <- survfit(Surv(time, status == 2) ~ 1, data = pbc)
plot(data) # Plots the result of survfit
plot(data, conf.int = "none") # New parameter
In order to try to make a reproducible example:
f <- function(x, y){
b <- x^2
c <- y^2
d <- x+y
return(list(one = b, two = c, three = d))
}
dat <- f(3, 2)
So using plot(dat)
I would like to get the same as plot(dat$one, dat$two)
. I would also like to add one more (new) parameter that could be set to TRUE
/FALSE
.
Is this possible?
Upvotes: 1
Views: 1089
Reputation: 35402
I think you might be looking for classes. You can use the S3 system for this.
For your survival
example, data
has the class survfit
(see class(data)
). Then using plot(data)
will look for a function called plot.survfit
. That is actually a non-exported function in the survival
package, at survival:::plot.survfit
.
You can easily do the same for your package. For example, have a function that creates an object of class my_class
, and then define a plotting method for that class:
f <- function(x, y){
b <- x^2
c <- y^2
d <- x+y
r <- list(one = b, two = c, three = d)
class(r) <- c('list', 'my_class') # this is the important bit.
r
}
plot.my_class <- function(x) {
plot(x$one, x$two)
}
Now your code should work:
dat <- f(3, 2)
plot(dat)
You can put anything in plot.my_class
you want, including additional arguments, as long as your first argument is x
and is the my_class
object.
plot
now calls plot.my_class
, since dat
is of class my_class
.
You can also add other methods, e.g. for print
.
There are many different plotting functions that can be called with plot
for different classes, see methods(plot)
Also see Hadley's Advanced R book chapter on S3.
Upvotes: 5