vec_slice2() prototype

vec_slice2.R

library(tidyverse)

vec_slice2 <- function(x, i) {
  if (!is.data.frame(x)) {
    return(.subset2(x, i))
  }
  # FIXME: Other special cases?
  row <- vctrs::vec_slice(x, i)
  lists <- map_lgl(row, ~ is.list(.x) && !is.data.frame(.x))
  length_one_elements <- map_lgl(
    row[lists],
    ~ {
      (is.atomic(.x[[1]]) || is.data.frame(.x[[1]])) && vctrs::vec_size(.x[[1]]) == 1
    }
  )

  out <- as.list(row)
  out[lists] <- map(out[lists], 1)

  # Edge case: nested columns of length one
  out[lists][length_one_elements] <- map(
    out[lists][length_one_elements],
    ~ {
      attr(.x, "scalar") <- TRUE
      .x
    }
  )

  out
}

tibble_row2 <- function(...) {
  data <- rlang::list2(...)
  atomic <- map_lgl(
    data,
    ~ (is.data.frame(.x) || (is.atomic(.x))) && vctrs::vec_size(.x) == 1 && is.null(attr(.x, "scalar"))
  )
  data[!atomic] <- map(data[!atomic], list)
  as_tibble(data)
}

input <- tibble(
  x = 1:3,
  y = list(1, 2:3, letters[4:6]),
  z = tibble(a = 4:6, b = 7:9),
  w = vctrs::list_of(tibble(c = 1L), tibble(c = 2:3), tibble(c = 4:6))
)

input
#> # A tibble: 3 × 4
#>       x y           z$a    $b                  w
#>   <int> <list>    <int> <int> <list<tibble[,1]>>
#> 1     1 <dbl [1]>     4     7            [1 × 1]
#> 2     2 <int [2]>     5     8            [2 × 1]
#> 3     3 <chr [3]>     6     9            [3 × 1]
vec_slice2(input, 1)
#> $x
#> [1] 1
#> 
#> $y
#> [1] 1
#> attr(,"scalar")
#> [1] TRUE
#> 
#> $z
#> # A tibble: 1 × 2
#>       a     b
#>   <int> <int>
#> 1     4     7
#> 
#> $w
#> # A tibble: 1 × 1
#>       c
#>   <int>
#> 1     1
vec_slice2(input, 2)
#> $x
#> [1] 2
#> 
#> $y
#> [1] 2 3
#> 
#> $z
#> # A tibble: 1 × 2
#>       a     b
#>   <int> <int>
#> 1     5     8
#> 
#> $w
#> # A tibble: 2 × 1
#>       c
#>   <int>
#> 1     2
#> 2     3
vec_slice2(input, 3)
#> $x
#> [1] 3
#> 
#> $y
#> [1] "d" "e" "f"
#> 
#> $z
#> # A tibble: 1 × 2
#>       a     b
#>   <int> <int>
#> 1     6     9
#> 
#> $w
#> # A tibble: 3 × 1
#>       c
#>   <int>
#> 1     4
#> 2     5
#> 3     6

roundtrip <- bind_rows(
  tibble_row2(!!!vec_slice2(input, 1)),
  tibble_row2(!!!vec_slice2(input, 2)),
  tibble_row2(!!!vec_slice2(input, 3)),
)

waldo::compare(input, roundtrip)
#> `attr(old$y[[1]], 'scalar')` is absent
#> `attr(new$y[[1]], 'scalar')` is a logical vector (TRUE)
#> 
#> `old$w` is an S3 object of class <vctrs_list_of/vctrs_vctr/list>, a list
#> `new$w` is a list