From f5c152ba6892705df83cd18dee499c079ec50a18 Mon Sep 17 00:00:00 2001
From: Benjamin Gruenbaum <benjamingr@gmail.com>
Date: Fri, 21 Jan 2022 18:42:21 +0200
Subject: [PATCH] stream: add drop and take

---
 lib/internal/streams/operators.js      | 23 +++++++++
 test/parallel/test-stream-drop-take.js | 67 ++++++++++++++++++++++++++
 2 files changed, 90 insertions(+)
 create mode 100644 test/parallel/test-stream-drop-take.js

diff --git a/lib/internal/streams/operators.js b/lib/internal/streams/operators.js
index 1218942b6414e4..accc91b7a5d657 100644
--- a/lib/internal/streams/operators.js
+++ b/lib/internal/streams/operators.js
@@ -236,10 +236,33 @@ async function* flatMap(fn, options) {
   }
 }
 
+async function* drop(number) {
+  validateInteger(number, 'number', 0);
+  for await (const val of this) {
+    if (number-- <= 0) {
+      yield await val;
+    }
+  }
+}
+
+
+async function* take(number) {
+  validateInteger(number, 'number', 0);
+  for await (const val of this) {
+    if (number-- > 0) {
+      yield val;
+    } else {
+      return;
+    }
+  }
+}
+
 module.exports.streamReturningOperators = {
+  drop,
   filter,
   flatMap,
   map,
+  take,
 };
 
 module.exports.promiseReturningOperators = {
diff --git a/test/parallel/test-stream-drop-take.js b/test/parallel/test-stream-drop-take.js
new file mode 100644
index 00000000000000..67265dad98cc15
--- /dev/null
+++ b/test/parallel/test-stream-drop-take.js
@@ -0,0 +1,67 @@
+'use strict';
+
+const common = require('../common');
+const {
+  Readable,
+} = require('stream');
+const { deepStrictEqual, rejects } = require('assert');
+
+const { from } = Readable;
+{
+  // Synchronous streams
+  (async () => {
+    deepStrictEqual(await from([1, 2, 3]).drop(2).toArray(), [3]);
+    deepStrictEqual(await from([1, 2, 3]).take(1).toArray(), [1]);
+    deepStrictEqual(await from([]).drop(2).toArray(), []);
+    deepStrictEqual(await from([]).take(1).toArray(), []);
+    deepStrictEqual(await from([1, 2, 3]).drop(1).take(1).toArray(), [2]);
+    deepStrictEqual(await from([1, 2]).drop(0).toArray(), [1, 2]);
+    deepStrictEqual(await from([1, 2]).take(0).toArray(), []);
+  })().then(common.mustCall());
+  // Asynchronous streams
+  (async () => {
+    const fromAsync = (...args) => from(...args).map(async (x) => x);
+    deepStrictEqual(await fromAsync([1, 2, 3]).drop(2).toArray(), [3]);
+    deepStrictEqual(await fromAsync([1, 2, 3]).take(1).toArray(), [1]);
+    deepStrictEqual(await fromAsync([]).drop(2).toArray(), []);
+    deepStrictEqual(await fromAsync([]).take(1).toArray(), []);
+    deepStrictEqual(await fromAsync([1, 2, 3]).drop(1).take(1).toArray(), [2]);
+    deepStrictEqual(await fromAsync([1, 2]).drop(0).toArray(), [1, 2]);
+    deepStrictEqual(await fromAsync([1, 2]).take(0).toArray(), []);
+  })().then(common.mustCall());
+  // Infinite streams
+  // Asynchronous streams
+  (async () => {
+    const naturals = () => from(async function*() {
+      let i = 1;
+      while (true) {
+        yield i++;
+      }
+    }());
+    deepStrictEqual(await naturals().take(1).toArray(), [1]);
+    deepStrictEqual(await naturals().drop(1).take(1).toArray(), [2]);
+    const next10 = [11, 12, 13, 14, 15, 16, 17, 18, 19, 20];
+    deepStrictEqual(await naturals().drop(10).take(10).toArray(), next10);
+    deepStrictEqual(await naturals().take(5).take(1).toArray(), [1]);
+  })().then(common.mustCall());
+}
+{
+  // Error cases
+  // TODO(benjamingr) check if the spec coerces
+  const invalidArgs = [
+    '5',
+    undefined,
+    null,
+    {},
+    [],
+    from([1, 2, 3]),
+    Promise.resolve(5),
+  ];
+
+  for (const example of invalidArgs) {
+    rejects(
+      from([]).take(example).toArray(),
+      /ERR_INVALID_ARG_TYPE/
+    ).then(common.mustCall());
+  }
+}