Allow n to be a recyclable vector in cal_advance(). See tidyverts/tsi…

…bble#18 for discussion.
DavisVaughan · Sep 2, 2018 · 1facc9e · 1facc9e
1 parent d9e7dcb
commit 1facc9e
Show file tree

Hide file tree

Showing 4 changed files with 37 additions and 7 deletions.
diff --git a/R/cal_advance.R b/R/cal_advance.R
@@ -1,3 +1,6 @@
+# TODO) Document that we allow dates and n to be vectors. Should unit and convention
+# be vectors as well?
+
 #' @export
 cal_advance <- function(dates, 
                         cal = cal_create(), 
@@ -13,6 +16,10 @@ cal_advance <- function(dates,
 
   n <- as_integer(n)
 
+  recycle_n <- max(length(dates), length(n))
+  dates <- recycle(dates, recycle_n)
+  n <- recycle(n, recycle_n)
+
   cal_advance_cpp(cal, dates, n, unit, convention, end_of_month)
 }
 

diff --git a/R/utils.R b/R/utils.R
@@ -29,3 +29,29 @@ empty_date <- function() {
   class(empty_numeric) <- "Date"
   empty_numeric
 }
+
+# This recycling does not allow for multiples (2 becomes 4 kind of thing)
+# It only recycles length 1 to length n, or does nothing if n_x == n already
+recycle <- function(x, n) {
+  x_n <- length(x)
+
+  valid_length <- (x_n == 1L) | (x_n == n) 
+
+  if(!valid_length) {
+    stop(
+      "Object to be recycled is length ", 
+      x_n, 
+      " but should be length 1 or ", 
+      n, 
+      ".",
+      call. = FALSE
+    )
+  }
+
+  # length 1, recycle
+  if(x_n == 1L) {
+    x <- rep(x, times = n)
+  }
+
+  x
+}
diff --git a/src/RcppExports.cpp b/src/RcppExports.cpp
@@ -20,14 +20,14 @@ BEGIN_RCPP
 END_RCPP
 }
 // cal_advance_cpp
-std::vector<QuantLib::Date> cal_advance_cpp(List cal, std::vector<QuantLib::Date> dates, const int& n, const std::string& unit, const std::string& convention, bool end_of_month);
+std::vector<QuantLib::Date> cal_advance_cpp(List cal, std::vector<QuantLib::Date> dates, Rcpp::IntegerVector n, const std::string& unit, const std::string& convention, bool end_of_month);
 RcppExport SEXP _calendarrr_cal_advance_cpp(SEXP calSEXP, SEXP datesSEXP, SEXP nSEXP, SEXP unitSEXP, SEXP conventionSEXP, SEXP end_of_monthSEXP) {
 BEGIN_RCPP
     Rcpp::RObject rcpp_result_gen;
     Rcpp::RNGScope rcpp_rngScope_gen;
     Rcpp::traits::input_parameter< List >::type cal(calSEXP);
     Rcpp::traits::input_parameter< std::vector<QuantLib::Date> >::type dates(datesSEXP);
-    Rcpp::traits::input_parameter< const int& >::type n(nSEXP);
+    Rcpp::traits::input_parameter< Rcpp::IntegerVector >::type n(nSEXP);
     Rcpp::traits::input_parameter< const std::string& >::type unit(unitSEXP);
     Rcpp::traits::input_parameter< const std::string& >::type convention(conventionSEXP);
     Rcpp::traits::input_parameter< bool >::type end_of_month(end_of_monthSEXP);

diff --git a/src/cal_advance.cpp b/src/cal_advance.cpp
@@ -19,7 +19,7 @@ using namespace Rcpp;
 // [[Rcpp::export]]
 std::vector<QuantLib::Date> cal_advance_cpp(List cal, 
                                             std::vector<QuantLib::Date> dates, 
-                                            const int &n,
+                                            Rcpp::IntegerVector n,
                                             const std::string &unit,
                                             const std::string &convention,
                                             bool end_of_month = false) {
@@ -30,9 +30,6 @@ std::vector<QuantLib::Date> cal_advance_cpp(List cal,
   // Create BusinessDayConvention object
   QuantLib::BusinessDayConvention bdc = get_business_day_convention(convention);
 
-  // Integer = QL_INTEGER = int
-  QuantLib::Integer n_ql_int = static_cast<QuantLib::Integer>(n);
-
   // Create TimeUnit
   QuantLib::TimeUnit time_unit = get_time_unit(unit);
 
@@ -41,7 +38,7 @@ std::vector<QuantLib::Date> cal_advance_cpp(List cal,
   std::vector<QuantLib::Date> advanced_dates(n_dates);
 
   for (int i = 0; i < n_dates; i++) {
-    advanced_dates[i] = cal_obj.advance(dates[i], n_ql_int, time_unit, bdc, end_of_month);
+    advanced_dates[i] = cal_obj.advance(dates[i], n[i], time_unit, bdc, end_of_month);
   }
 
   // Reset