From 63cf732ab48c8fdb93285478c2786d0f288bbf2f Mon Sep 17 00:00:00 2001
From: Jarred Sumner <jarred@jarredsumner.com>
Date: Fri, 2 Aug 2024 23:05:48 -0700
Subject: [PATCH] Support async iterators in fs.promises.writeFile (#13044)

---
 src/js/node/fs.promises.ts                    | 76 +++++++++++++++++++
 ...-promises-writeFile-async-iterator.test.ts | 53 +++++++++++++
 2 files changed, 129 insertions(+)
 create mode 100644 test/js/node/fs/fs-promises-writeFile-async-iterator.test.ts

diff --git a/src/js/node/fs.promises.ts b/src/js/node/fs.promises.ts
index ada7640b74..5960cc01bf 100644
--- a/src/js/node/fs.promises.ts
+++ b/src/js/node/fs.promises.ts
@@ -200,6 +200,15 @@ const exports = {
   },
   writeFile: function (fileHandleOrFdOrPath, ...args) {
     fileHandleOrFdOrPath = fileHandleOrFdOrPath?.[kFd] ?? fileHandleOrFdOrPath;
+    if (
+      !$isTypedArrayView(args[0]) &&
+      typeof args[0] !== "string" &&
+      ($isCallable(args[0]?.[Symbol.iterator]) || $isCallable(args[0]?.[Symbol.asyncIterator]))
+    ) {
+      $debug("fs.promises.writeFile async iterator slow path!");
+      // Node accepts an arbitrary async iterator here
+      return writeFileAsyncIterator(fileHandleOrFdOrPath, ...args);
+    }
     return _writeFile(fileHandleOrFdOrPath, ...args);
   },
   readlink: fs.readlink.bind(fs),
@@ -570,3 +579,70 @@ function throwEBADFIfNecessary(fn, fd) {
     throw err;
   }
 }
+
+async function writeFileAsyncIteratorInner(fd, iterable, encoding) {
+  const writer = Bun.file(fd).writer();
+
+  const mustRencode = !(encoding === "utf8" || encoding === "utf-8" || encoding === "binary" || encoding === "buffer");
+  let totalBytesWritten = 0;
+
+  try {
+    for await (let chunk of iterable) {
+      if (mustRencode && typeof chunk === "string") {
+        $debug("Re-encoding chunk to", encoding);
+        chunk = Buffer.from(chunk, encoding);
+      }
+
+      const prom = writer.write(chunk);
+      if (prom && $isPromise(prom)) {
+        totalBytesWritten += await prom;
+      } else {
+        totalBytesWritten += prom;
+      }
+    }
+  } finally {
+    await writer.end();
+  }
+
+  return totalBytesWritten;
+}
+
+async function writeFileAsyncIterator(fdOrPath, iterable, optionsOrEncoding, flag, mode) {
+  let encoding;
+  if (typeof optionsOrEncoding === "object") {
+    encoding = optionsOrEncoding?.encoding ?? (encoding || "utf8");
+    flag = optionsOrEncoding?.flag ?? (flag || "w");
+    mode = optionsOrEncoding?.mode ?? (mode || 0o666);
+  } else if (typeof optionsOrEncoding === "string" || optionsOrEncoding == null) {
+    encoding = optionsOrEncoding || "utf8";
+    flag ??= "w";
+    mode ??= 0o666;
+  }
+
+  if (!Buffer.isEncoding(encoding)) {
+    // ERR_INVALID_OPT_VALUE_ENCODING was removed in Node v15.
+    throw new TypeError(`Unknown encoding: ${encoding}`);
+  }
+
+  let mustClose = typeof fdOrPath === "string";
+  if (mustClose) {
+    // Rely on fs.open for further argument validaiton.
+    fdOrPath = await fs.open(fdOrPath, flag, mode);
+  }
+
+  let totalBytesWritten = 0;
+
+  try {
+    totalBytesWritten = await writeFileAsyncIteratorInner(fdOrPath, iterable, encoding);
+  } finally {
+    if (mustClose) {
+      try {
+        if (typeof flag === "string" && !flag.includes("a")) {
+          await fs.ftruncate(fdOrPath, totalBytesWritten);
+        }
+      } finally {
+        await fs.close(fdOrPath);
+      }
+    }
+  }
+}
diff --git a/test/js/node/fs/fs-promises-writeFile-async-iterator.test.ts b/test/js/node/fs/fs-promises-writeFile-async-iterator.test.ts
new file mode 100644
index 0000000000..ad26caca5c
--- /dev/null
+++ b/test/js/node/fs/fs-promises-writeFile-async-iterator.test.ts
@@ -0,0 +1,53 @@
+import { test, expect, mock } from "bun:test";
+import { writeFile } from "fs/promises";
+import { tempDirWithFiles } from "harness";
+test("fs.promises.writeFile async iterator", async () => {
+  const dir = tempDirWithFiles("fs-promises-writeFile-async-iterator", {
+    "file1.txt": "0 Hello, world!",
+  });
+  const path = dir + "/file2.txt";
+
+  const stream = async function* () {
+    yield "1 ";
+    yield "Hello, ";
+    yield "world!";
+  };
+
+  await writeFile(path, stream());
+  expect(await Bun.file(path).text()).toBe("1 Hello, world!");
+
+  const bufStream = async function* () {
+    yield Buffer.from("2 ");
+    yield Buffer.from("Hello, ");
+    yield Buffer.from("world!");
+  };
+
+  await writeFile(path, bufStream());
+
+  expect(await Bun.file(path).text()).toBe("2 Hello, world!");
+});
+
+test("fs.promises.writeFile async iterator throws on invalid input", async () => {
+  const dir = tempDirWithFiles("fs-promises-writeFile-async-iterator", {
+    "file1.txt": "0 Hello, world!",
+  });
+  const symbolStream = async function* () {
+    yield Symbol("lolwhat");
+  };
+
+  expect(() => writeFile(dir + "/file2.txt", symbolStream())).toThrow();
+  expect(() =>
+    writeFile(
+      dir + "/file3.txt",
+      (async function* () {
+        yield "once";
+        throw new Error("good");
+      })(),
+    ),
+  ).toThrow("good");
+  const fn = {
+    [Symbol.asyncIterator]: mock(() => {}),
+  };
+  expect(() => writeFile(dir, fn)).toThrow();
+  expect(fn[Symbol.asyncIterator]).not.toBeCalled();
+});