Skip to content

Instantly share code, notes, and snippets.

@robertzk
Last active Apr 7, 2016
Embed
What would you like to do?
R-Python itertools rosetta stone

A translation of Python's itertools to R.

count(10) # 10 11 12 13 14 ...
count <- function(x, step = 1) {
  function() { (x <<- x + step) }
}
cycle([1,2,3]) # [1, 2, 3, 1, 2, 3, ...]
cycle <- function(x) {
 iterator <- 0
 function() {
   iterator <<- iterator + 1
   if (iterator == length(x) + 1) {
      iterator <<- 1
   }
   x[[iterator]]
 }
}
> cycle(c(1,2,3)) -> fn
> fn()
[1] 1
> fn()
[1] 2
> fn()
[1] 3
> fn()
[1] 1
> fn()
[1] 2
> fn()
[1] 3
repeat(10, 3) # 10 10 10
repeat2 <- function(elem, n = Inf) {
  count <- 0
  function() { 
    count <<- count + 1
    if (count <= n) elem else NULL
  }
}
chain([1,2], [3,4]) # 1 2 3 4
chain <- function(...) {
  args <- list(...)
  iterator <- 0
  args_iterator <- 1
  function() {
    iterator <<- iterator + 1
    compare_len <- length(args[[args_iterator]])
    while (iterator > compare_len) {
      iterator <<- 1
      args_iterator <<- args_iterator + 1
    }
    if (length(args) < args_iterator) {
      NULL
    } else {
      args[[args_iterator]][[iterator]]
    }
  }
}
compress('ABCDEF', [1,0,1,0,1,1]) # A C E F
compress <- function(x, y) {
  x[as.logical(y)]
}
compress(c("A","B","C","D","E","F"), c(1,0,1,0,1,1))
dropwhile(lambda x: x<5, [1,4,6,4,1]) # 6 4 1
takewhile <- function(pred, seq) {
  `%|%` <- function(x, y) if (is.na(x)) y else x
  seq[seq_len(Position(Negate(pred), seq) %|% length(seq))]
}
[list(g) for v, g in groupby([1,1,2,4,6,3,5,2,7,9], lambda x: x % 2)]
# [[1, 1], [2, 4, 6], [3, 5], [2], [7, 9]]
groupby <- function(iterable, keyfunc) {
  keys <- lapply(iterable, keyfunc)
  matches <- c(TRUE, as.logical(Map(identical, keys[-length(keys)], keys[-1L])))
  unname(split(iterable, cumsum(1 - matches)))
}
# groupby(list(5,8,12,14,5,27,25, 31), function(x) x %/% 10)
# [[5,8],[12,14],[5],[27,25],[31]]
ifilter(lambda x: x%2, range(10)) # 1 3 5 7 9
ifilter <- Filter
ifilterfalse(lambda x: x%2, range(10)) # 0 2 4 6 8
ifilterfalse <- function(f, x) Filter(Negate(f), x)
islice('ABCDEFG', 2, None) # C D E F G
islice <- function(seq, start, stop = NULL, step = 1) {
  if (is.null(stop) || stop > length(seq)) stop <- length(seq)
  if (stop < start) seq[logical(0)]
  else seq[start + seq_len((stop - start + 1 + step) %/% step) * step - step]
}
# > islice(1:10, 2, 8, 3)
# [1] 2 5 8
imap(pow, (2,3,10), (5,2,3)) # 32 9 1000
imap <- Map
# > as.integer(Map(`^`, c(2,3,10), c(5,2,3)))
# [1]   32    9 1000
starmap(pow, [(2,5), (3,2), (10,3)]) # 32 9 1000
flip <- function(fn) function(x, y) fn(y, x)
starmap <- function(func, seq) {
  lapply(seq, flip(do.call), func)
}
tee <- function(it, n) {
  lapply(seq_len(n), function(.) it
}
takewhile(lambda x: x<5, [1,4,6,4,1]) # 1 4
takewhile <- function(pred, seq) {
  `%|%` <- function(x, y) if (is.na(x)) y else x
  seq[seq_len(Position(pred, seq) %|% length(seq))]
}
izip('ABCD', 'xy') # Ax By
izip <- function(...) {
  n <- seq_len(min(vapply(list(...), length, integer(1))))
  unname(do.call(Map, c(list(c), lapply(list(...), function(.) .[n]))))
}
# > izip(LETTERS[1:4], c("x","y"))
# [[1]]
# [1] "A" "x"
# 
# [[2]]
# [1] "B" "y"
izip_longest # izip_longest('ABCD', 'xy', fillvalue='-') --> Ax By C- D-
izip_longest <- function(..., fillvalue = NULL) {
  n <- max(vapply(list(...), length, integer(1)))
  output <- lapply(list(...), function(.) c(., list(fillvalue)[rep(1, n - length(.))]))
  unname(do.call(Map, c(list(c), output)))
}
# > izip_longest(LETTERS[1:4], c("x","y"), c("FOO", "BAR"), fillvalue = '-')
# [[1]]
# [1] "A"   "x"   "FOO"
#
# [[2]]
# [1] "B"   "y"   "BAR"
#
# [[3]]
# [1] "C" "-" "-"
#
# [[4]]
# [1] "D" "-" "-"

product, combinations (hint: utils::combn), and permutations left as an exercise to the reader.

And just for fun:

flip <- function(f) {
  force(f)
  function(...) { 
    dots <- eval(substitute(alist(...)))
    dots[c(1, 2)] <- dots[c(2, 1)]
    do.call(f, dots)
  }
}
# > flip(`/`)(3, 4)
# [1] 1.333333
@ithayer
Copy link

ithayer commented Apr 7, 2016

rather than be inspired by itertools, you should be inspired by clojure

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment