Support async generators as body

szmarczak · szmarczak · commit 854430f01350 · 2021-07-14T02:35:55.000+02:00
Fixes #1735
diff --git a/documentation/3-streams.md b/documentation/3-streams.md
@@ -110,6 +110,31 @@ An object representing how much data have been downloaded.
 
 An object representing how much data have been uploaded.
 
+**Note:**
+> - When a chunk is greater than `highWaterMark`, the progress won't be emitted. The body needs to be split into chunks.
+
+```js
+import got from 'got';
+
+const body = Buffer.alloc(1024 * 1024); // 1MB
+
+function* chunkify(buffer, chunkSize = 64 * 1024) {
+	for (let pos = 0; pos < buffer.byteLength; pos += chunkSize) {
+		yield buffer.subarray(pos, pos + chunkSize)
+	}
+}
+
+const stream = got.stream.post('https://httpbin.org/anything', {
+	body: chunkify(body)
+});
+
+stream.resume();
+
+stream.on('uploadProgress', progress => {
+	console.log(progress);
+});
+```
+
 ### `stream.timings`
 
 **Type: [`Timings`](typescript.md#timings)**
diff --git a/documentation/4-pagination.md b/documentation/4-pagination.md
@@ -188,7 +188,7 @@ The reason `filter` looks exactly the same like `shouldContinue` is that the lat
 The `filter` function is needed as well, because in the same response we can get results with different timestamps.
 
 ```js
-import got from '../../dist/source/index.js';
+import got from 'got';
 import Bourne from '@hapi/bourne';
 
 const max = Date.now() - 1000 * 86400 * 7;
diff --git a/documentation/examples/h2c.js b/documentation/examples/h2c.js
@@ -1,5 +1,5 @@
 import http2 from 'http2-wrapper';
-import got from 'got';
+import got from '../../dist/source/index.js';
 
 let sessions = {};
 const getSession = ({origin}) => {
diff --git a/source/core/index.ts b/source/core/index.ts
@@ -918,13 +918,38 @@ export default class Request extends Duplex implements RequestEvents<Request> {
 		this.emit('request', request);
 	}
 
+	private async _asyncWrite(chunk: any): Promise<void> {
+		return new Promise((resolve, reject) => {
+			super.write(chunk, error => {
+				if (error) {
+					reject(error);
+					return;
+				}
+
+				resolve();
+			});
+		});
+	}
+
 	private _sendBody() {
 		// Send body
 		const {body} = this.options;
 		const currentRequest = this.redirectUrls.length === 0 ? this : this._request ?? this;
 
 		if (is.nodeStream(body)) {
 			body.pipe(currentRequest);
+		} else if (is.generator(body)) {
+			(async () => {
+				try {
+					for await (const chunk of body) {
+						await this._asyncWrite(chunk);
+					}
+
+					super.end();
+				} catch (error) {
+					this._beforeError(error);
+				}
+			})();
 		} else {
 			this._unlockWrite();
 
diff --git a/source/core/options.ts b/source/core/options.ts
@@ -845,6 +845,11 @@ export default class Options {
 
 		this._merging = true;
 
+		// Always merge `isStream` first
+		if ('isStream' in options) {
+			this.isStream = options.isStream!;
+		}
+
 		try {
 			let push = false;
 
@@ -1087,12 +1092,12 @@ export default class Options {
 
 	Since Got 12, the `content-length` is not automatically set when `body` is a `fs.createReadStream`.
 	*/
-	get body(): string | Buffer | Readable | undefined {
+	get body(): string | Buffer | Readable | Generator | AsyncGenerator | undefined {
 		return this._internals.body;
 	}
 
-	set body(value: string | Buffer | Readable | undefined) {
-		assert.any([is.string, is.buffer, is.nodeStream, is.undefined], value);
+	set body(value: string | Buffer | Readable | Generator | AsyncGenerator | undefined) {
+		assert.any([is.string, is.buffer, is.nodeStream, is.generator, is.asyncGenerator, is.undefined], value);
 
 		if (is.nodeStream(value)) {
 			assert.truthy(value.readable);
