From 0ac4ee6ea36c4906e1d039d371a553ef5cb00947 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andr=C3=A9=20Ver=C3=ADssimo?=
 <211358+averissimo@users.noreply.github.com>
Date: Fri, 3 Nov 2023 15:54:22 +0100
Subject: [PATCH] Add [[ and [[<- custom methods for JoinKeys

---
 NAMESPACE                       |   2 +
 R/join_keys.R                   | 221 +++++++++++++++++++++++++-------
 inst/WORDLIST                   |   1 +
 man/assert_join_keys.Rd         |  26 +++-
 man/assert_join_keys_alike.Rd   |  29 +++++
 man/join_keys.Rd                |  22 ++++
 man/merge_join_keys.Rd          |   1 +
 tests/testthat/test-join_keys.R |   2 +-
 tests/testthat/test-parents.R   |   5 +
 9 files changed, 256 insertions(+), 53 deletions(-)
 create mode 100644 man/assert_join_keys_alike.Rd

diff --git a/NAMESPACE b/NAMESPACE
index 2a1ff8aa7..49df5e6ad 100644
--- a/NAMESPACE
+++ b/NAMESPACE
@@ -2,6 +2,8 @@
 
 S3method("[",JoinKeys)
 S3method("[<-",JoinKeys)
+S3method("[[",JoinKeys)
+S3method("[[<-",JoinKeys)
 S3method("join_keys<-",JoinKeys)
 S3method("join_keys<-",teal_data)
 S3method("parents<-",JoinKeys)
diff --git a/R/join_keys.R b/R/join_keys.R
index 40a9afd20..f13fee7ce 100644
--- a/R/join_keys.R
+++ b/R/join_keys.R
@@ -156,6 +156,14 @@ join_keys <- function(...) {
 #' jk["ds1"]
 #' jk[["ds1"]]
 `[.JoinKeys` <- function(join_keys_obj, dataset_1 = NULL, dataset_2 = NULL) {
+  if (checkmate::test_integerish(dataset_1)) {
+    return(NextMethod("[", join_keys_obj))
+  } else if (length(dataset_1) > 1) {
+    res <- lapply(dataset_1, function(x) get_join_key(join_keys_obj, x, dataset_2))
+    names(res) <- dataset_1
+    class(res) <- class(new_join_keys())
+    return(res)
+  }
   get_join_key(join_keys_obj, dataset_1, dataset_2)
 }
 
@@ -182,9 +190,105 @@ join_keys <- function(...) {
 #' jk["ds1"] <- "primary_key"
 #' jk
 `[<-.JoinKeys` <- function(join_keys_obj, dataset_1, dataset_2 = dataset_1, value) {
+  if (checkmate::test_integerish(dataset_1)) {
+    stop(paste(
+      "Assigment via index number is not supported with JoinKeys object,",
+      "please use a dataset name as index and one at a time."
+    ))
+  } else if (length(dataset_1) > 1) {
+    stop(paste(
+      "Assigment of multiple JoinKeys at the same time is not supported,",
+      "please do one at a time."
+    ))
+  }
   add_key(join_keys_obj, dataset_1, dataset_2, value)
 }
 
+#' @rdname join_keys
+#' @export
+#' @examples
+#'
+#' jk <- join_keys(join_key("ds1", "ds2", "key"))
+#' jk[["ds1"]]
+#' jk[["ds1", "ds2"]]
+`[[.JoinKeys` <- function(join_keys_obj, dataset_1 = NULL, dataset_2 = NULL, value) {
+  if (!is.null(dataset_1) && !is.null(dataset_2)) {
+    return(join_keys_obj[[dataset_1]][[dataset_2]])
+  }
+  NextMethod("[[", jk)
+}
+
+#' @rdname join_keys
+#' @export
+#' @examples
+#'
+#' jk <- join_keys()
+#' jk[["ds1"]] <- list()
+#' jk[["ds2"]][["ds3"]] <- "key"
+#' jk[["ds3", "ds4"]] <- "new_key"
+#'
+#' jk <- join_keys()
+#' jk[["ds1"]] <- list()
+#' jk[["ds2"]][["ds3"]] <- "key"
+#' jk[["ds4"]] <- list(ds5 = "new")
+#' jk[["ds6", "ds7"]] <- "yada"
+#' jk[["ds8", "ds9"]] <- c(A = "B", "C")
+`[[<-.JoinKeys` <- function(join_keys_obj, dataset_1 = NULL, dataset_2 = NULL, value) {
+  checkmate::assert_string(dataset_1)
+  checkmate::assert_string(dataset_2, null.ok = TRUE)
+
+  # Accepting 2 subscripts
+  if (!is.null(dataset_2)) {
+    checkmate::assert_character(value)
+    # Normalize value
+    new_join_key <- join_key(dataset_1, dataset_2, value)
+    dataset_1 <- new_join_key$dataset_1
+    dataset_2 <- new_join_key$dataset_2
+    value <- new_join_key$keys
+
+    if (is.null(join_keys_obj[[dataset_1]])) {
+      join_keys_obj[[dataset_1]] <- list()
+    }
+    join_keys_obj[[dataset_1]][[dataset_2]] <- value
+    return(join_keys_obj)
+  }
+
+  # Accepting 1 subscript with valid `value` formal
+  checkmate::assert_list(value, names = "named", types = "character", null.ok = TRUE)
+
+  join_keys_obj <- NextMethod("[[<-", jk)
+
+  # Keep original parameters as variables will be overwritten for `NextMethod` call
+  original_value <- value
+  ds1 <- dataset_1
+
+  # Iterate on all new values to create symmetrical pair
+  for (ds2 in names(value)) {
+    if (ds2 == ds1) next
+
+    value <- rlang::`%||%`(join_keys_obj[[ds2]], list())
+    new_value <- original_value[[ds2]]
+    if (
+      checkmate::test_character(new_value, min.len = 1) &&
+        all(is.null(names(new_value)))
+    ) {
+      new_value <- setNames(new_value, new_value)
+    } else if (
+      checkmate::test_character(new_value, min.len = 1)
+    ) {
+      # Invert key
+      new_value <- setNames(names(new_value), new_value)
+    }
+
+    # Change variables for NextMethod call
+    dataset_1 <- ds2
+    value[[ds1]] <- new_value
+    join_keys_obj <- NextMethod("[[<-", join_keys_obj)
+  }
+
+  join_keys_obj
+}
+
 #' Mutate `JoinKeys` with a new values
 #'
 #' @description `r lifecycle::badge("experimental")`
@@ -329,6 +433,7 @@ merge_join_keys.default <- function(join_keys_obj, new_join_keys) {
 #' jk2["ds1", "ds3"] <- "new_col"
 #'
 #' merge_join_keys(jk1, jk2)
+#' merge_join_keys(jk1, list(jk2))
 merge_join_keys.JoinKeys <- function(join_keys_obj, new_join_keys) {
   assert_join_keys(join_keys_obj)
 
@@ -336,6 +441,8 @@ merge_join_keys.JoinKeys <- function(join_keys_obj, new_join_keys) {
     new_join_keys <- list(new_join_keys)
   }
 
+  lapply(new_join_keys, assert_join_keys_alike)
+
   checkmate::assert_list(new_join_keys, types = c("JoinKeys"), min.len = 1)
 
   result <- join_keys_obj
@@ -388,9 +495,10 @@ print.JoinKeys <- function(x, ...) {
 #'
 #' @keywords internal
 new_join_keys <- function() {
-  result <- list()
-  class(result) <- c("JoinKeys", "list")
-  result
+  structure(
+    list(),
+    class = c("JoinKeys", "list")
+  )
 }
 
 #' Get value of a single relationship pair
@@ -407,11 +515,6 @@ new_join_keys <- function() {
 #'
 #' @keywords internal
 get_join_key <- function(join_keys_obj, dataset_1 = NULL, dataset_2 = NULL) {
-  if (checkmate::test_integerish(dataset_1, len = 2)) {
-    # if dataset_1 is an index integet vector, then return itself
-    #  trick to make the following work: join_keys(jk)["ds1", "ds2"] <- "key"
-    return(join_keys_obj)
-  }
   checkmate::assert_string(dataset_1, null.ok = TRUE)
   if (missing(dataset_2)) {
     # protection if dataset_2 is passed through by a function
@@ -446,39 +549,9 @@ get_join_key <- function(join_keys_obj, dataset_1 = NULL, dataset_2 = NULL) {
 add_key <- function(join_keys_obj, dataset_1, dataset_2 = dataset_1, value) {
   checkmate::assert_string(dataset_1)
   checkmate::assert_string(dataset_2, null.ok = TRUE)
+  checkmate::assert_character(value)
 
-  # Normalize value
-  new_join_key <- join_key(dataset_1, dataset_2, value)
-  dataset_1 <- new_join_key$dataset_1
-  dataset_2 <- new_join_key$dataset_2
-  value <- new_join_key$keys
-
-  # Create pair ds_1 -> ds_2
-  if (is.null(join_keys_obj[[dataset_1]])) join_keys_obj[[dataset_1]] <- list()
-
-  join_keys_obj[[dataset_1]][[dataset_2]] <- value
-
-  # Primary key, do nothing else
-  if (identical(dataset_1, dataset_2)) {
-    return(join_keys_obj)
-  }
-
-  # Create symmetrical pair ds_2 -> ds_1
-  if (is.null(join_keys_obj[[dataset_2]])) join_keys_obj[[dataset_2]] <- list()
-
-  if (
-    checkmate::test_character(value, min.len = 1) &&
-      all(is.null(names(value)))
-  ) {
-    value <- setNames(names(value), value)
-  } else if (
-    checkmate::test_character(value, min.len = 1)
-  ) {
-    # Invert key
-    value <- setNames(names(value), value)
-  }
-
-  join_keys_obj[[dataset_2]][[dataset_1]] <- value
+  join_keys_obj[[dataset_1, dataset_2]] <- value
   join_keys_obj
 }
 
@@ -500,27 +573,81 @@ join_pair <- function(join_keys_obj, join_key_obj) {
   join_keys_obj
 }
 
-#' Check the JoinKeys class membership of an argument
+#' Assert the JoinKeys class membership of an argument
 #' @inheritParams checkmate::assert_class
-#' @param extra_classes (`character` vector) with extra classes to check. Can be used
 #'
 #' @return `x` invisibly
 #'
 #' @keywords internal
-assert_join_keys <- function(x, .var.name = checkmate::vname(x)) {
-  checkmate::assert_class(x, classes = c("JoinKeys"), .var.name = .var.name)
+assert_join_keys <- function(x, .var.name = checkmate::vname(x), add = NULL) {
+  if (missing(x)) {
+    stop(sprintf("argument \"%s\" is missing, with no default", .var.name))
+  }
+
+  res <- check_join_keys(x)
+  checkmate::makeAssertion(x, res, var.name = .var.name, add)
+}
+
+#' @rdname assert_join_keys_alike
+#' @examples
+check_join_keys <- function(x) {
+  checkmate::check_class(x, classes = c("JoinKeys", "list"))
 }
 
 #' @rdname assert_join_keys
 #' @keywords internal
 test_join_keys <- function(x) {
-  checkmate::test_class(x, classes = c("JoinKeys"))
+  checkmate::makeTest(check_join_keys(x))
 }
 
 #' @rdname assert_join_keys
 #' @keywords internal
-expect_join_keys <- function(x) {
-  checkmate::expect_class(x, classes = c("JoinKeys"))
+expect_join_keys <- function(x, info = NULL, label = vname(x)) {
+  checkmate::makeExpectation(x, check_join_keys(x), info = info, label = label)
+}
+
+#' Assert the JoinKeys class membership of an argument
+#' @inheritParams checkmate::assert_class
+#'
+#' @return `x` invisibly
+#'
+#' @keywords internal
+assert_join_keys_alike <- function(x, .var.name = checkmate::vname(x), add = NULL) {
+  if (missing(x)) {
+    stop(sprintf("argument \"%s\" is missing, with no default", .var.name))
+  }
+  res <- check_join_keys_alike(x)
+
+  checkmate::makeAssertion(x, res, var.name = .var.name, add)
+}
+
+#' @rdname assert_join_keys
+#' @examples
+#' check_join_keys_alike(list("ds1" = list("key")))
+#' check_join_keys_alike(list("ds1" = list(ds2 = "key")))
+check_join_keys_alike <- function(x) {
+  result <- checkmate::check_list(x, names = "named", types = "list")
+  if (checkmate::test_string(result)) {
+    return(result)
+  }
+  result <- all(
+    vapply(
+      x,
+      function(el) {
+        checkmate::test_list(el, types = "character", names = "named")
+      },
+      logical(1)
+    )
+  )
+  if (isFALSE(all(result))) {
+    return(
+      paste(
+        "Elements of list may only be named lists with a vector of  type `character`",
+        "(that may be named or partially named)"
+      )
+    )
+  }
+  result
 }
 
 #' Helper function to assert if two key sets contain incompatible keys
diff --git a/inst/WORDLIST b/inst/WORDLIST
index b1ea8634f..e33bd1f7e 100644
--- a/inst/WORDLIST
+++ b/inst/WORDLIST
@@ -14,4 +14,5 @@ Reproducibility
 reproducibility
 returens
 SCDA
+testthat
 UI
diff --git a/man/assert_join_keys.Rd b/man/assert_join_keys.Rd
index d073530f8..512ae9025 100644
--- a/man/assert_join_keys.Rd
+++ b/man/assert_join_keys.Rd
@@ -4,13 +4,16 @@
 \alias{assert_join_keys}
 \alias{test_join_keys}
 \alias{expect_join_keys}
-\title{Check the JoinKeys class membership of an argument}
+\alias{check_join_keys_alike}
+\title{Assert the JoinKeys class membership of an argument}
 \usage{
-assert_join_keys(x, .var.name = checkmate::vname(x))
+assert_join_keys(x, .var.name = checkmate::vname(x), add = NULL)
 
 test_join_keys(x)
 
-expect_join_keys(x)
+expect_join_keys(x, info = NULL, label = vname(x))
+
+check_join_keys_alike(x)
 }
 \arguments{
 \item{x}{[any]\cr
@@ -20,12 +23,25 @@ Object to check.}
 Name of the checked object to print in assertions. Defaults to
 the heuristic implemented in \code{\link[checkmate]{vname}}.}
 
-\item{extra_classes}{(\code{character} vector) with extra classes to check. Can be used}
+\item{add}{[\code{AssertCollection}]\cr
+Collection to store assertion messages. See \code{\link[checkmate]{AssertCollection}}.}
+
+\item{info}{[\code{character(1)}]\cr
+Extra information to be included in the message for the testthat reporter.
+See \code{\link[testthat]{expect_that}}.}
+
+\item{label}{[\code{character(1)}]\cr
+Name of the checked object to print in messages. Defaults to
+the heuristic implemented in \code{\link[checkmate]{vname}}.}
 }
 \value{
 \code{x} invisibly
 }
 \description{
-Check the JoinKeys class membership of an argument
+Assert the JoinKeys class membership of an argument
+}
+\examples{
+check_join_keys_alike(list("ds1" = list("key")))
+check_join_keys_alike(list("ds1" = list(ds2 = "key")))
 }
 \keyword{internal}
diff --git a/man/assert_join_keys_alike.Rd b/man/assert_join_keys_alike.Rd
new file mode 100644
index 000000000..e4ce72646
--- /dev/null
+++ b/man/assert_join_keys_alike.Rd
@@ -0,0 +1,29 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/join_keys.R
+\name{check_join_keys}
+\alias{check_join_keys}
+\alias{assert_join_keys_alike}
+\title{Assert the JoinKeys class membership of an argument}
+\usage{
+check_join_keys(x)
+
+assert_join_keys_alike(x, .var.name = checkmate::vname(x), add = NULL)
+}
+\arguments{
+\item{x}{[any]\cr
+Object to check.}
+
+\item{.var.name}{[\code{character(1)}]\cr
+Name of the checked object to print in assertions. Defaults to
+the heuristic implemented in \code{\link[checkmate]{vname}}.}
+
+\item{add}{[\code{AssertCollection}]\cr
+Collection to store assertion messages. See \code{\link[checkmate]{AssertCollection}}.}
+}
+\value{
+\code{x} invisibly
+}
+\description{
+Assert the JoinKeys class membership of an argument
+}
+\keyword{internal}
diff --git a/man/join_keys.Rd b/man/join_keys.Rd
index eb1cf7c72..11455c31c 100644
--- a/man/join_keys.Rd
+++ b/man/join_keys.Rd
@@ -8,6 +8,8 @@
 \alias{join_keys<-.teal_data}
 \alias{[.JoinKeys}
 \alias{[<-.JoinKeys}
+\alias{[[.JoinKeys}
+\alias{[[<-.JoinKeys}
 \title{Create a \code{JoinKeys} out of a list of \code{JoinKeySet} objects}
 \usage{
 cdisc_join_keys(...)
@@ -23,6 +25,10 @@ join_keys(join_keys_obj) <- value
 \method{[}{JoinKeys}(join_keys_obj, dataset_1 = NULL, dataset_2 = NULL)
 
 \method{[}{JoinKeys}(join_keys_obj, dataset_1, dataset_2 = dataset_1) <- value
+
+\method{[[}{JoinKeys}(join_keys_obj, dataset_1 = NULL, dataset_2 = NULL, value)
+
+\method{[[}{JoinKeys}(join_keys_obj, dataset_1 = NULL, dataset_2 = NULL) <- value
 }
 \arguments{
 \item{...}{optional, a \code{JoinKeySet} objects created using the \code{join_key} function.}
@@ -111,4 +117,20 @@ jk["ds1", "ds2"] <- "(new) pair key"
 # Creates primary key by only defining `dataset_1`
 jk["ds1"] <- "primary_key"
 jk
+
+jk <- join_keys(join_key("ds1", "ds2", "key"))
+jk[["ds1"]]
+jk[["ds1", "ds2"]]
+
+jk <- join_keys()
+jk[["ds1"]] <- list()
+jk[["ds2"]][["ds3"]] <- "key"
+jk[["ds3", "ds4"]] <- "new_key"
+
+jk <- join_keys()
+jk[["ds1"]] <- list()
+jk[["ds2"]][["ds3"]] <- "key"
+jk[["ds4"]] <- list(ds5 = "new")
+jk[["ds6", "ds7"]] <- "yada"
+jk[["ds8", "ds9"]] <- c(A = "B", "C")
 }
diff --git a/man/merge_join_keys.Rd b/man/merge_join_keys.Rd
index 5cadc5ec5..4d96d2b7a 100644
--- a/man/merge_join_keys.Rd
+++ b/man/merge_join_keys.Rd
@@ -31,4 +31,5 @@ jk2 <- join_keys()
 jk2["ds1", "ds3"] <- "new_col"
 
 merge_join_keys(jk1, jk2)
+merge_join_keys(jk1, list(jk2))
 }
diff --git a/tests/testthat/test-join_keys.R b/tests/testthat/test-join_keys.R
index bd27355ac..93a91b9fc 100644
--- a/tests/testthat/test-join_keys.R
+++ b/tests/testthat/test-join_keys.R
@@ -556,7 +556,7 @@ testthat::test_that("merge_join_keys merges mutually exclusive data", {
   testthat::expect_true(all(join_keys(y) %in% join_keys(z)) && all(join_keys(z) %in% join_keys(y)))
   testthat::expect_true(all(join_keys(y) %in% join_keys(x)) && all(join_keys(x) %in% join_keys(y)))
 
-  testthat::expect_identical(names(join_keys(z)), c("A", "B", "Z", "Y"))
+  testthat::expect_identical(names(z), c("A", "B", "Z", "Y"))
   testthat::expect_equal(length(join_keys(z)), 4)
   testthat::expect_identical(join_keys(z)$A$B, c("a" = "b"))
   testthat::expect_identical(join_keys(z)$B$A, c("b" = "a"))
diff --git a/tests/testthat/test-parents.R b/tests/testthat/test-parents.R
index 080847371..0805bff9a 100644
--- a/tests/testthat/test-parents.R
+++ b/tests/testthat/test-parents.R
@@ -189,7 +189,10 @@ testthat::test_that("assert_parent_child throws error if no join_keys exist for
   join_keys(jk) <- list(
     join_key("df1", "df1", c("id" = "id"))
   )
+  # Change class as trick to allow for corrupt JoinKeys
+  class(jk) <- "list"
   jk[["df2"]][["df1"]] <- "id"
+  class(jk) <- class(new_join_keys())
   parents(jk) <- list(df1 = character(0), df2 = "df1", df3 = "df1")
   testthat::expect_error(
     assert_parent_child(jk),
@@ -202,7 +205,9 @@ testthat::test_that("assert_parent_child throws error if no join_keys exist for
   join_keys(jk) <- list(
     join_key("df1", "df1", c("id" = "id"))
   )
+  class(jk) <- "list"
   jk[["df1"]][["df2"]] <- "id"
+  class(jk) <- class(new_join_keys())
   parents(jk) <- list(df1 = character(0), df2 = "df1", df3 = "df1")
   testthat::expect_error(
     assert_parent_child(jk),