str_replace should ignore NA when using function as replacement

DESCRIPTION

@@ -19,7 +19,7 @@ Depends:
 Imports: 
     glue,
     magrittr,
-    stringi (>= 0.4.1)
+    stringi (>= 1.1.6)
 Suggests: 
     covr,
     htmltools,

R/replace.r

@@ -159,7 +159,7 @@ str_replace_na <- function(string, replacement = "NA") {
 
 str_transform <- function(string, pattern, replacement) {
   loc <- str_locate(string, pattern)
-  str_sub(string, loc) <- replacement(str_sub(string, loc))
+  str_sub(string, loc, omit_na = TRUE) <- replacement(str_sub(string, loc))
   string
 }
 str_transform_all <- function(string, pattern, replacement) {

R/sub.r

@@ -15,6 +15,8 @@
 #'   matrix to `start`.
 #'
 #'   Negative values count backwards from the last character.
+#' @param omit_na single logical value; if \code{TRUE}, missing values in any of the arguments
+#' provided will result in an unchanged input; replacement function only
 #' @param value replacement string
 #' @return A character vector of substring from `start` to `end`
 #'   (inclusive). Will be length of longest input argument.
@@ -50,6 +52,15 @@
 #' str_sub(x, -1, -1) <- "K"; x
 #' str_sub(x, -2, -2) <- "GHIJ"; x
 #' str_sub(x, 2, -2) <- ""; x
+#'
+#' # If you want to keep the original if some argument is NA,
+#' # use omit_na = TRUE
+#' x1 <- x2 <- x3 <- x4 <- "AAA"
+#' str_sub(x1, 1, NA) <- "B"
+#' str_sub(x2, 1, 2) <- NA
+#' str_sub(x3, 1, NA, omit_na = TRUE) <- "B"
+#' str_sub(x4, 1, 2, omit_na = TRUE) <- NA
+#' x1; x2; x3; x4
 str_sub <- function(string, start = 1L, end = -1L) {
   if (is.matrix(start)) {
     stri_sub(string, from = start)
@@ -61,11 +72,11 @@ str_sub <- function(string, start = 1L, end = -1L) {
 
 #' @export
 #' @rdname str_sub
-"str_sub<-" <- function(string, start = 1L, end = -1L, value) {
+"str_sub<-" <- function(string, start = 1L, end = -1L, omit_na = FALSE,  value) {
   if (is.matrix(start)) {
-    stri_sub(string, from = start) <- value
+    stri_sub(string, from = start, omit_na = omit_na) <- value
   } else {
-    stri_sub(string, from = start, to = end) <- value
+    stri_sub(string, from = start, to = end, omit_na = omit_na) <- value
   }
   string
 }

tests/testthat/test-replace.r

@@ -69,6 +69,10 @@ test_that("replacement can be different length", {
   expect_equal(str_replace_all("abc", "a|c", double), "aabcc")
 })
 
+test_that("replacement with NA works", {
+  expect_equal(str_replace("abc", "z", toupper), "abc")
+})
+
 # fix_replacement ---------------------------------------------------------
 
 test_that("$ are escaped", {

tests/testthat/test-sub.r

@@ -69,6 +69,15 @@ test_that("replacement works", {
 
   str_sub(x, 2, -2) <- ""
   expect_equal(x, "AH")
+})
 
+test_that("replacement with NA works", {
+  x <- "BBCDEF"
+  str_sub(x, NA) <- "A"
+  expect_equal(x, NA_character_)
 
+  x <- "BBCDEF"
+  str_sub(x, NA, omit_na = TRUE) <- "A"
+  str_sub(x, 1, 1, omit_na = TRUE) <- NA
+  expect_equal(x, "BBCDEF")
 })