fs: writeFile support AsyncIterable, Iterable & Stream as data argument

HiroyukiYagihashi · HiroyukiYagihashi · commit 3e95264c3797 · 2021-02-24T09:21:10.000+09:00
Fixes: #37391
diff --git a/doc/api/fs.md b/doc/api/fs.md
@@ -3833,6 +3833,9 @@ details.
 <!-- YAML
 added: v0.1.29
 changes:
+  - version: REPLACEME
+    pr-url: https://github.com/nodejs/node/pull/37490
+    description: `data` argument support `AsyncIterable`, `Iterable` & `Stream`.
   - version: v15.2.0
     pr-url: https://github.com/nodejs/node/pull/35993
     description: The options argument may include an AbortSignal to abort an
@@ -3866,7 +3869,8 @@ changes:
 -->
 
 * `file` {string|Buffer|URL|integer} filename or file descriptor
-* `data` {string|Buffer|TypedArray|DataView|Object}
+* `data` {string|Buffer|TypedArray|DataView|Object|AsyncIterable|Iterable
+  |Stream}
 * `options` {Object|string}
   * `encoding` {string|null} **Default:** `'utf8'`
   * `mode` {integer} **Default:** `0o666`
diff --git a/lib/internal/fs/promises.js b/lib/internal/fs/promises.js
@@ -10,6 +10,7 @@ const kReadFileMaxChunkSize = 2 ** 14;
 const kWriteFileMaxChunkSize = 2 ** 14;
 
 const {
+  ArrayIsArray,
   ArrayPrototypePush,
   Error,
   MathMax,
@@ -21,6 +22,8 @@ const {
   PromiseResolve,
   SafeArrayIterator,
   Symbol,
+  SymbolAsyncIterator,
+  SymbolIterator,
   Uint8Array,
 } = primordials;
 
@@ -41,7 +44,7 @@ const {
   ERR_INVALID_ARG_VALUE,
   ERR_METHOD_NOT_IMPLEMENTED,
 } = codes;
-const { isArrayBufferView } = require('internal/util/types');
+const { isArrayBufferView, isTypedArray } = require('internal/util/types');
 const { rimrafPromises } = require('internal/fs/rimraf');
 const {
   copyObject,
@@ -273,7 +276,21 @@ async function fsCall(fn, handle, ...args) {
 }
 
 async function writeFileHandle(filehandle, data, signal) {
-  // `data` could be any kind of typed array.
+  if (signal?.aborted) {
+    throw lazyDOMException('The operation was aborted', 'AbortError');
+  }
+  if (isIterable(data)) {
+    for await (const buf of data) {
+      if (signal?.aborted) {
+        throw lazyDOMException('The operation was aborted', 'AbortError');
+      }
+      await filehandle.write(buf);
+      if (signal?.aborted) {
+        throw lazyDOMException('The operation was aborted', 'AbortError');
+      }
+    }
+    return;
+  }
   data = new Uint8Array(data.buffer, data.byteOffset, data.byteLength);
   let remaining = data.length;
   if (remaining === 0) return;
@@ -663,7 +680,7 @@ async function writeFile(path, data, options) {
   options = getOptions(options, { encoding: 'utf8', mode: 0o666, flag: 'w' });
   const flag = options.flag || 'w';
 
-  if (!isArrayBufferView(data)) {
+  if (!isArrayBufferView(data) && !isIterable(data)) {
     validateStringAfterArrayBufferView(data, 'data');
     data = Buffer.from(data, options.encoding || 'utf8');
   }
@@ -678,6 +695,18 @@ async function writeFile(path, data, options) {
   return PromisePrototypeFinally(writeFileHandle(fd, data), fd.close);
 }
 
+function isIterable(obj) {
+  if (obj == null) {
+    return false;
+  }
+
+  return SymbolAsyncIterator in obj || (
+    SymbolIterator in obj &&
+    typeof obj !== 'string' &&
+    !ArrayIsArray(obj) &&
+    !isTypedArray(obj));
+}
+
 async function appendFile(path, data, options) {
   options = getOptions(options, { encoding: 'utf8', mode: 0o666, flag: 'a' });
   options = copyObject(options);
diff --git a/test/parallel/test-fs-promises-writefile.js b/test/parallel/test-fs-promises-writefile.js
@@ -7,20 +7,70 @@ const path = require('path');
 const tmpdir = require('../common/tmpdir');
 const assert = require('assert');
 const tmpDir = tmpdir.path;
+const { Readable } = require('stream');
 
 tmpdir.refresh();
 
 const dest = path.resolve(tmpDir, 'tmp.txt');
 const otherDest = path.resolve(tmpDir, 'tmp-2.txt');
 const buffer = Buffer.from('abc'.repeat(1000));
 const buffer2 = Buffer.from('xyz'.repeat(1000));
+const stream = Readable.from(['a', 'b', 'c']);
+const stream2 = Readable.from(['a', 'b', 'c']);
+const iterable = {
+  [Symbol.iterator]: function*() {
+    yield 'a';
+    yield 'b';
+    yield 'c';
+  }
+};
+const asyncIterable = {
+  async* [Symbol.asyncIterator]() {
+    yield 'a';
+    yield 'b';
+    yield 'c';
+  }
+};
 
 async function doWrite() {
   await fsPromises.writeFile(dest, buffer);
   const data = fs.readFileSync(dest);
   assert.deepStrictEqual(data, buffer);
 }
 
+async function doWriteStream() {
+  await fsPromises.writeFile(dest, stream);
+  let expected = '';
+  for await (const v of stream2) expected += v;
+  const data = fs.readFileSync(dest, 'utf-8');
+  assert.deepStrictEqual(data, expected);
+}
+
+async function doWriteStreamWithCancel() {
+  const controller = new AbortController();
+  const { signal } = controller;
+  process.nextTick(() => controller.abort());
+  assert.rejects(fsPromises.writeFile(otherDest, stream, { signal }), {
+    name: 'AbortError'
+  });
+}
+
+async function doWriteIterable() {
+  await fsPromises.writeFile(dest, iterable);
+  let expected = '';
+  for await (const v of iterable) expected += v;
+  const data = fs.readFileSync(dest, 'utf-8');
+  assert.deepStrictEqual(data, expected);
+}
+
+async function doWriteAsyncIterable() {
+  await fsPromises.writeFile(dest, asyncIterable);
+  let expected = '';
+  for await (const v of asyncIterable) expected += v;
+  const data = fs.readFileSync(dest, 'utf-8');
+  assert.deepStrictEqual(data, expected);
+}
+
 async function doWriteWithCancel() {
   const controller = new AbortController();
   const { signal } = controller;
@@ -55,4 +105,8 @@ doWrite()
   .then(doAppend)
   .then(doRead)
   .then(doReadWithEncoding)
+  .then(doWriteStream)
+  .then(doWriteStreamWithCancel)
+  .then(doWriteIterable)
+  .then(doWriteAsyncIterable)
   .then(common.mustCall());
