Add slice function

caolan · Mar 19, 2015 · 7233df5 · 7233df5
1 parent 0003a6a
commit 7233df5
Show file tree

Hide file tree

Showing 2 changed files with 83 additions and 39 deletions.
diff --git a/lib/index.js b/lib/index.js
@@ -2326,45 +2326,60 @@ Stream.prototype.split = function () {
 exposeMethod('split');
 
 /**
- * Creates a new Stream with the first `n` values from the source.
+ * Creates a new Stream with the values from the source in the range of `start` to `end`.
  *
- * @id take
+ * @id slice
  * @section Transforms
- * @name Stream.take(n)
- * @param {Number} n - integer representing number of values to read from source
+ * @name Stream.slice(start, end)
+ * @param {Number} start - integer representing index to start reading from source
+ * @param {Number} stop - integer representing index to stop reading from source
  * @api public
  *
- * _([1, 2, 3, 4]).take(2) // => 1, 2
+ * _([1, 2, 3, 4]).slice(1, 3) // => 2, 3
  */
+Stream.prototype.slice = function(start, end) {
+    var index = 0;
+    start = typeof start != 'number' ? 0 : start;
+    end = typeof end != 'number' ? Infinity : end;
 
-Stream.prototype.take = function (n) {
-    if (n === 0) {
+    if (start < 0 || start >= end) {
         return _([]);
     }
     var s = this.consume(function (err, x, push, next) {
+        var done = x === nil;
         if (err) {
             push(err);
-            if (n > 0) {
-                next();
-            }
-            else {
-                push(null, nil);
-            }
         }
-        else if (x === nil) {
-            push(null, nil);
+        else if (!done && index++ >= start) {
+            push(null, x);
+        }
+
+        if (!done && index < end) {
+            next();
         }
         else {
-            n--;
-            push(null, x);
-            if (n > 0) {
-                next();
-            }
-            else {
-                push(null, nil);
-            }
+            push(null, nil);
         }
     });
+    s.id = 'slice:' + s.id;
+    return s;
+};
+exposeMethod('slice');
+
+/**
+ * Creates a new Stream with the first `n` values from the source.
+ *
+ * @id take
+ * @section Transforms
+ * @name Stream.take(n)
+ * @param {Number} n - integer representing number of values to read from source
+ * @api public
+ *
+ * _([1, 2, 3, 4]).take(2) // => 1, 2
+ */
+
+Stream.prototype.take = function (n) {
+    var s = this.slice(0, n);
     s.id = 'take:' + s.id;
     return s;
 };
@@ -2388,22 +2403,7 @@ Stream.prototype.drop = function (n) {
     if (n <= 0) {
         return this;
     }
-    return this.consume(function (err, x, push, next) {
-        if (err) {
-            push(err);
-            next();
-        }
-        else if (x === nil) {
-            push(null, nil);
-        }
-        else {
-            n--;
-            if (n < 0) {
-                push(null, x);
-            }
-            next();
-        }
-    });
+    return this.slice(n);
 };
 exposeMethod('drop');
 

diff --git a/test/test.js b/test/test.js
@@ -949,6 +949,50 @@ exports['take'] = function (test) {
     test.done();
 };
 
+exports['slice'] = {
+    setUp: function (cb) {
+        this.input = [1, 2, 3, 4, 5];
+        this.expected = [3, 4, 5];
+        this.tester = function (expected, test) {
+            return function (xs) {
+                test.same(xs, expected);
+            };
+        };
+        cb();
+    },
+    'arrayStream': function (test) {
+        test.expect(5);
+        _(this.input).slice(2, 6).toArray(this.tester([3, 4, 5], test));
+        _(this.input).slice(2).toArray(this.tester([3, 4, 5], test));
+        _(this.input).slice().toArray(this.tester([1, 2, 3, 4, 5], test));
+        _(this.input).slice(-1, 6).toArray(this.tester([], test));
+        _(this.input).slice(0).toArray(this.tester([1, 2, 3, 4, 5], test));
+        test.done();
+    },
+    'partial application': function (test) {
+        test.expect(1);
+        var s = _(this.input);
+        _.slice(1, 4)(s).toArray(this.tester([2, 3, 4], test));
+        test.done();
+    },
+    'error': function (test) {
+        test.expect(2);
+        var s = _(function (push, next) {
+            push(null, 1),
+            push(new Error('Slice error')),
+            push(null, 2),
+            push(null, 3),
+            push(null, 4),
+            push(null, 5),
+            push(null, _.nil)
+        });
+        s.slice(2, 4).errors(errorEquals(test, 'Slice error'))
+            .toArray(this.tester([3, 4], test));
+        test.done();
+    },
+    'noValueOnError': noValueOnErrorTest(_.slice(2, 3))
+};
+
 exports['take - noValueOnError'] = noValueOnErrorTest(_.take(1));
 
 exports['take - errors'] = function (test) {