feat: Provide takeShort() & takeWithFallback()

koraa · koraa · commit bafa834ff48c · 2019-08-27T19:41:04.000+02:00
diff --git a/src/sequence.js b/src/sequence.js
@@ -1168,6 +1168,7 @@ const skipWhile = curry('skipWhile', (seq, pred) => {
  * elements if the input sequence was shorter than `no` elements.
  *
  * @function
+ * @alias takeShort
  * @param {Sequence} seq Any sequence for which iter() is defined
  * @param {Number} no The number of elements to take
  * @returns {Iterator} The first element for which pred returns false
@@ -1179,10 +1180,12 @@ const tryTake = curry('tryTake', (seq, no) => {
     range0(no),
     map(() => tryNext(seq, Nothing)),
     takeUntilVal(Nothing),
-    list
+    list,
   );
 });
 
+const takeShort = tryTake;
+
 /**
  * Version of tryTake that will throw IteratorEnded
  * if the given iterable is too short.
@@ -1196,7 +1199,24 @@ const tryTake = curry('tryTake', (seq, no) => {
 const take = curry('take', (seq, no) => pipe(
   range0(no),
   map(() => next(seq)),
-  list
+  list,
+));
+
+/**
+ * Yields an iterator of the first `no` elements in the given
+ * in the sequence. If the sequence is too short, the fallback
+ * parameter will be substituted.
+ *
+ * @function
+ * @param {Sequence} seq Any sequence for which iter() is defined
+ * @param {Number} no The number of elements to take
+ * @param {Any} fallback The value to supply if the input sequence is too short
+ * @returns {Array} The elements taken
+ */
+const takeWithFallback = curry('takeWithFallback', (seq, no, fallback) => pipe(
+  concat(range0(no)),
+  map(() => tryNext(seq, fallback)),
+  list,
 ));
 
 /**
@@ -1676,7 +1696,9 @@ module.exports = {
   skip,
   skipWhile,
   tryTake,
+  takeShort,
   take,
+  takeWithFallback,
   takeWhile,
   takeUntilVal,
   takeDef,
diff --git a/test/sequence.test.js b/test/sequence.test.js
@@ -22,6 +22,7 @@ const {
   trySecond, tryLast, seqEq, each, find, tryFind, contains, count, list,
   uniq, join, dict, obj, into, foldl, foldr, any, all, sum, product, map,
   filter, reject, reverse, enumerate, trySkip, skip, skipWhile, tryTake,
+  takeShort, takeWithFallback,
   take, takeWhile, takeUntilVal, takeDef, flat, concat, prepend, append,
   mapSort, zipLeast, zip, zipLongest, zipLeast2, zip2, zipLongest2,
   slidingWindow, trySlidingWindow, lookahead, mod, union, union2,
@@ -309,18 +310,26 @@ it('skipWhile', () => {
 it('take/...', () => {
   ckEqSeq(tryTake(4)(range0(10)), [0, 1, 2, 3]);
   ckEqSeq(tryTake(4)(range0(2)), [0, 1]);
+  ckEqSeq(tryTake(4)([]), []);
+  ckEqSeq(takeWithFallback(4, 99)(range0(2)), [0, 1, 99, 99]);
+  ckEqSeq(takeWithFallback(4, 99)([]), [99, 99, 99, 99]);
   ckEqSeq(take(2)(range0(2)), [0, 1]);
   ckThrows(IteratorEnded, () => take(4)(range0(2)));
 
-  const it = iter(range0(8));
+  const it = iter(range0(12));
   ckEq(take(it, 0), []);
-  ckEq(take(it, 3), [0,1,2]);
-  ckEq(take(it, 2), [3,4]);
+  ckEq(take(it, 3), [0, 1, 2]);
+  ckEq(take(it, 2), [3, 4]);
   ckEq(tryTake(it, 0), []);
-  ckEq(tryTake(it, 1), [5]);
-  ckEq(tryTake(it, 2), [6,7]);
+  ckEq(takeShort(it, 1), [5]);
+  ckEq(tryTake(it, 2), [6, 7]);
+  ckEq(takeWithFallback(it, 0, null), []);
+  ckEq(takeWithFallback(it, 2, null), [8, 9]);
+  ckEq(takeWithFallback(it, 4, null), [10, 11, null, null]);
+  ckEq(takeWithFallback(it, 0, null), []);
   ckEq(take(it, 0), []);
-  ckEq(tryTake(it, 0), []);
+  ckEq(takeShort(it, 0), []);
+  ckEq(takeWithFallback(it, 0, null), []);
 });
 
 it('takeWhile()', () => {
