fs: writeFile support AsyncIterable, Iterable & Stream as `da…

…ta` argument Fixes: nodejs#37391
yagipy · Feb 21, 2021 · 3274cad · 3274cad
1 parent 75cc41e
commit 3274cad
Show file tree

Hide file tree

Showing 3 changed files with 76 additions and 10 deletions.
diff --git a/doc/api/fs.md b/doc/api/fs.md
@@ -3866,7 +3866,7 @@ 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
@@ -663,19 +663,42 @@ async function writeFile(path, data, options) {
   options = getOptions(options, { encoding: 'utf8', mode: 0o666, flag: 'w' });
   const flag = options.flag || 'w';
 
-  if (!isArrayBufferView(data)) {
-    validateStringAfterArrayBufferView(data, 'data');
-    data = Buffer.from(data, options.encoding || 'utf8');
+  if (options.signal?.aborted) {
+    throw lazyDOMException('The operation was aborted', 'AbortError');
   }
 
-  if (path instanceof FileHandle)
-    return writeFileHandle(path, data, options.signal);
+  if (isStream(data) || isIterable(data)) {
+    const fd = await open(path, flag, options.mode);
 
-  const fd = await open(path, flag, options.mode);
-  if (options.signal?.aborted) {
-    throw lazyDOMException('The operation was aborted', 'AbortError');
+    try {
+      for await (const buf of data) {
+        await fd.write(buf);
+      }
+    } finally {
+      await fd.close();
+    }
+  } else {
+    if (!isArrayBufferView(data)) {
+      validateStringAfterArrayBufferView(data, 'data');
+      data = Buffer.from(data, options.encoding || 'utf8');
+    }
+
+    if (path instanceof FileHandle) {
+      return writeFileHandle(path, data, options.signal);
+    }
+
+    const fd = await open(path, flag, options.mode);
+    return PromisePrototypeFinally(writeFileHandle(fd, data), fd.close);
   }
-  return PromisePrototypeFinally(writeFileHandle(fd, data), fd.close);
+}
+
+function isStream(obj) {
+  return !!(obj && typeof obj.pipe === 'function');
+}
+
+function isIterable(obj) {
+  return !!obj && typeof obj !== 'string' && !obj.buffer
+      && (typeof obj[Symbol.iterator] === 'function' || typeof obj[Symbol.asyncIterator] === 'function');
 }
 
 async function appendFile(path, data, options) {

diff --git a/test/parallel/test-fs-promises-writefile.js b/test/parallel/test-fs-promises-writefile.js
@@ -7,13 +7,29 @@ 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(["abc".repeat(1000)]);
+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);
@@ -50,9 +66,36 @@ async function doReadWithEncoding() {
   assert.deepStrictEqual(data, syncData);
 }
 
+async function doWriteStream() {
+  await fsPromises.writeFile(dest, stream);
+  let result = "";
+  for await (const v of stream) result += v;
+  const data = fs.readFileSync(dest);
+  assert.deepStrictEqual(data, result);
+}
+
+async function doWriteIterable() {
+  await fsPromises.writeFile(dest, iterable);
+  let result = "";
+  for await (const v of iterable) result += v;
+  const data = fs.readFileSync(dest);
+  assert.deepStrictEqual(data, result);
+}
+
+async function doWriteAsyncIterable() {
+  await fsPromises.writeFile(dest, asyncIterable);
+  let result = "";
+  for await (const v of iterable) result += v;
+  const data = fs.readFileSync(dest);
+  assert.deepStrictEqual(data, result);
+}
+
 doWrite()
   .then(doWriteWithCancel)
   .then(doAppend)
   .then(doRead)
   .then(doReadWithEncoding)
+  .then(doWriteStream)
+  .then(doWriteIterable)
+  .then(doWriteAsyncIterable)
   .then(common.mustCall());