diff --git a/Cargo.lock b/Cargo.lock index 2ae16d7b..e1124403 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1325,6 +1325,7 @@ name = "impit-node" version = "0.0.0" dependencies = [ "bytes", + "futures-util", "h2", "impit", "napi", diff --git a/impit-node/Cargo.toml b/impit-node/Cargo.toml index 26c57945..8df73f1f 100644 --- a/impit-node/Cargo.toml +++ b/impit-node/Cargo.toml @@ -16,6 +16,7 @@ h2 = "0.4.7" reqwest = { version="0.13.1" } tokio-stream = "0.1.17" bytes = "1.11.1" +futures-util = "0.3" [build-dependencies] napi-build = "2.3.2" diff --git a/impit-node/index.wrapper.js b/impit-node/index.wrapper.js index 3e9ed22e..91b74fd4 100644 --- a/impit-node/index.wrapper.js +++ b/impit-node/index.wrapper.js @@ -49,15 +49,22 @@ function toUint8Array(chunk) { return typeof chunk === 'string' ? new TextEncoder().encode(chunk) : new Uint8Array(chunk); } -function concatUint8Arrays(chunks) { - const totalLength = chunks.reduce((acc, chunk) => acc + chunk.length, 0); - const result = new Uint8Array(totalLength); - let offset = 0; - for (const chunk of chunks) { - result.set(chunk, offset); - offset += chunk.length; - } - return result; +// Streamed bodies are passed to the native layer as a ReadableStream of bytes, so that the chunks +// reach the connection as they are produced instead of being buffered before the request starts. +function toByteStream(source) { + const reader = source instanceof ReadableStream ? source.getReader() : source[Symbol.asyncIterator](); + const next = reader.read?.bind(reader) ?? reader.next.bind(reader); + + return new ReadableStream({ + async pull(controller) { + const { done, value } = await next(); + if (done) { + controller.close(); + } else { + controller.enqueue(toUint8Array(value)); + } + }, + }); } function shouldRewriteRedirectToGet(httpStatus, method) { @@ -193,23 +200,9 @@ class Impit extends native.Impit { return { body: new Uint8Array(await body.arrayBuffer()), type: body.type }; } else if (body instanceof FormData) { return await this.#generateMultipartFormData(body); - } else if (body instanceof ReadableStream) { - const reader = body.getReader(); - const chunks = []; - let done = false; - while (!done) { - const { done: streamDone, value } = await reader.read(); - done = streamDone; - if (value != null) chunks.push(toUint8Array(value)); - } - return { body: concatUint8Arrays(chunks), type: '' }; - } else if (body && typeof body[Symbol.asyncIterator] === 'function') { - // Node.js streams (e.g. Readable.from(...)) and other async iterables. - const chunks = []; - for await (const chunk of body) { - chunks.push(toUint8Array(chunk)); - } - return { body: concatUint8Arrays(chunks), type: '' }; + } else if (body instanceof ReadableStream || typeof body[Symbol.asyncIterator] === 'function') { + // Web streams, Node.js streams (e.g. Readable.from(...)) and other async iterables. + return { bodyStream: toByteStream(body), type: '' }; } return { body, type: '' }; } @@ -238,8 +231,9 @@ class Impit extends native.Impit { options.headers = canonicalizeHeaders(options?.headers); if (options?.body) { - const { body: requestBody, type } = await this.#serializeBody(options.body); + const { body: requestBody, bodyStream, type } = await this.#serializeBody(options.body); options.body = requestBody; + options.bodyStream = bodyStream; if (type && !options.headers.some(([key]) => key.toLowerCase() === 'content-type')) { options.headers.push(['Content-Type', type]); } @@ -252,6 +246,7 @@ class Impit extends native.Impit { method: options.method, headers: options.headers, body: options.body, + bodyStream: options.bodyStream, timeout: options.timeout, forceHttp3: options.forceHttp3, signal: options.signal, @@ -316,6 +311,7 @@ class Impit extends native.Impit { method, headers, body: method === 'GET' ? undefined : options.body, + bodyStream: method === 'GET' ? undefined : options.bodyStream, }); const originalResponse = await Promise.race([ @@ -349,6 +345,10 @@ class Impit extends native.Impit { url = new URL(location, url).toString(); method = shouldRewriteRedirectToGet(originalResponse.status, method) ? 'GET' : method; + if (options.bodyStream && method !== 'GET') { + throw new TypeError(`Cannot follow a redirect that resends a streaming request body: ${url}`); + } + continue; } } diff --git a/impit-node/src/lib.rs b/impit-node/src/lib.rs index 90ec326d..269694ba 100644 --- a/impit-node/src/lib.rs +++ b/impit-node/src/lib.rs @@ -1,6 +1,10 @@ use std::time::Duration; -use impit::{errors::ImpitError, impit::Impit, request::RequestOptions}; +use impit::{ + errors::ImpitError, + impit::Impit, + request::{ImpitBody, RequestOptions}, +}; use napi::{bindgen_prelude::ObjectFinalize, Env}; use napi_derive::napi; @@ -136,9 +140,10 @@ impl ImpitWrapper { .as_ref() .and_then(|init| init.method.to_owned()) .unwrap_or_default(); - let body = request_init - .and_then(|init| init.body) - .map(|array| array.to_vec().into()); + let body = request_init.and_then(|init| match (init.body, init.body_stream) { + (_, Some(stream)) => Some(ImpitBody::from_stream(stream.into_bytes())), + (bytes, None) => bytes.map(|array| array.to_vec().into()), + }); let response = if matches!(method, HttpMethod::Get | HttpMethod::Head) && body.is_some() { Err(ImpitError::BindingPassthroughError( diff --git a/impit-node/src/request.rs b/impit-node/src/request.rs index 59657050..e1a843ab 100644 --- a/impit-node/src/request.rs +++ b/impit-node/src/request.rs @@ -1,7 +1,25 @@ -use napi::bindgen_prelude::Uint8Array; +use bytes::Bytes; +use futures_util::{Stream, TryStreamExt}; +use napi::bindgen_prelude::{sys, FromNapiValue, ReadableStream, Reader, Uint8Array}; use napi_derive::napi; +/// A JS `ReadableStream` of byte chunks, read chunk by chunk as the request body is sent. +pub struct BodyStream(Reader); + +impl FromNapiValue for BodyStream { + unsafe fn from_napi_value(env: sys::napi_env, value: sys::napi_value) -> napi::Result { + let stream = unsafe { ReadableStream::::from_napi_value(env, value)? }; + Ok(Self(stream.read()?)) + } +} + +impl BodyStream { + pub fn into_bytes(self) -> impl Stream> { + self.0.map_ok(|chunk| Bytes::copy_from_slice(&chunk)) + } +} + #[derive(Default, Clone)] #[napi(string_enum = "UPPERCASE")] pub enum HttpMethod { @@ -24,7 +42,7 @@ pub enum HttpMethod { /// /// See {@link Impit.fetch} for usage. #[derive(Default)] -#[napi(object)] +#[napi(object, object_to_js = false)] pub struct RequestInit { /// HTTP method to use for the request. Default is `GET`. /// @@ -47,6 +65,9 @@ pub struct RequestInit { )] /// Request body. Can be a string, Buffer, ArrayBuffer, TypedArray, DataView, Blob, File, URLSearchParams, FormData or ReadableStream. pub body: Option, + /// Set by the JS wrapper instead of `body` when the body is a stream. Takes precedence over `body`. + #[napi(skip_typescript)] + pub body_stream: Option, /// Request timeout in milliseconds. Overrides the Impit-wide timeout option from {@link ImpitOptions.timeout}. pub timeout: Option, /// Force the request to use HTTP/3. If the server doesn't expect HTTP/3 or the Impit instance doesn't have HTTP/3 enabled (via the {@link ImpitOptions.http3} option), the request will fail. diff --git a/impit-node/test/basics.test.ts b/impit-node/test/basics.test.ts index f009f8cc..521b38ed 100644 --- a/impit-node/test/basics.test.ts +++ b/impit-node/test/basics.test.ts @@ -522,6 +522,20 @@ describe.each([ expect(json.data).toEqual(STRING_PAYLOAD); }); + // https://github.com/apify/impit/issues/513 + test('streamed bodies are sent chunk by chunk', async () => { + const chunks = ['{"Impit-Test":', '"foořžš"}']; + const response = await impit.fetch(localPostUrl, { + method: HttpMethod.Post, + body: Readable.from(chunks) as any, + }); + const json = await response.json(); + + expect(json.headers?.['transfer-encoding']).toBe('chunked'); + expect(json.headers?.['content-length']).toBeUndefined(); + expect(json.data).toEqual(STRING_PAYLOAD); + }); + test('Request with body preserves its Content-Type', async () => { const request = new Request(localPostUrl, { method: 'POST', diff --git a/impit-python/python/impit/impit.pyi b/impit-python/python/impit/impit.pyi index f1667a53..ce39e96d 100644 --- a/impit-python/python/impit/impit.pyi +++ b/impit-python/python/impit/impit.pyi @@ -6,7 +6,7 @@ from .headers import Headers from . import Browser from typing import Any -from collections.abc import Iterator, AsyncIterator +from collections.abc import AsyncIterable, AsyncIterator, Iterable, Iterator from contextlib import AbstractAsyncContextManager, AbstractContextManager @@ -518,7 +518,7 @@ class Client: def get( self, url: str, - content: bytes | bytearray | list[int] | None = None, + content: bytes | bytearray | list[int] | Iterable[bytes] | None = None, data: dict[str, str] | None = None, headers: dict[str, str] | None = None, timeout: float | str | None = USE_CLIENT_DEFAULT, @@ -538,7 +538,7 @@ class Client: def post( self, url: str, - content: bytes | bytearray | list[int] | None = None, + content: bytes | bytearray | list[int] | Iterable[bytes] | None = None, data: dict[str, str] | None = None, headers: dict[str, str] | None = None, timeout: float | str | None = USE_CLIENT_DEFAULT, @@ -559,7 +559,7 @@ class Client: def put( self, url: str, - content: bytes | bytearray | list[int] | None = None, + content: bytes | bytearray | list[int] | Iterable[bytes] | None = None, data: dict[str, str] | None = None, headers: dict[str, str] | None = None, timeout: float | str | None = USE_CLIENT_DEFAULT, @@ -579,7 +579,7 @@ class Client: def patch( self, url: str, - content: bytes | bytearray | list[int] | None = None, + content: bytes | bytearray | list[int] | Iterable[bytes] | None = None, data: dict[str, str] | None = None, headers: dict[str, str] | None = None, timeout: float | str | None = USE_CLIENT_DEFAULT, @@ -599,7 +599,7 @@ class Client: def delete( self, url: str, - content: bytes | bytearray | list[int] | None = None, + content: bytes | bytearray | list[int] | Iterable[bytes] | None = None, data: dict[str, str] | None = None, headers: dict[str, str] | None = None, timeout: float | str | None = USE_CLIENT_DEFAULT, @@ -619,7 +619,7 @@ class Client: def head( self, url: str, - content: bytes | bytearray | list[int] | None = None, + content: bytes | bytearray | list[int] | Iterable[bytes] | None = None, data: dict[str, str] | None = None, headers: dict[str, str] | None = None, timeout: float | str | None = USE_CLIENT_DEFAULT, @@ -639,7 +639,7 @@ class Client: def options( self, url: str, - content: bytes | bytearray | list[int] | None = None, + content: bytes | bytearray | list[int] | Iterable[bytes] | None = None, data: dict[str, str] | None = None, headers: dict[str, str] | None = None, timeout: float | str | None = USE_CLIENT_DEFAULT, @@ -659,7 +659,7 @@ class Client: def trace( self, url: str, - content: bytes | bytearray | list[int] | None = None, + content: bytes | bytearray | list[int] | Iterable[bytes] | None = None, data: dict[str, str] | None = None, headers: dict[str, str] | None = None, timeout: float | str | None = USE_CLIENT_DEFAULT, @@ -680,7 +680,7 @@ class Client: self, method: str, url: str, - content: bytes | bytearray | list[int] | None = None, + content: bytes | bytearray | list[int] | Iterable[bytes] | None = None, data: dict[str, str] | None = None, headers: dict[str, str] | None = None, timeout: float | str | None = USE_CLIENT_DEFAULT, @@ -704,7 +704,7 @@ class Client: self, method: str, url: str, - content: bytes | bytearray | list[int] | None = None, + content: bytes | bytearray | list[int] | Iterable[bytes] | None = None, data: dict[str, str] | None = None, headers: dict[str, str] | None = None, timeout: float | str | None = USE_CLIENT_DEFAULT, @@ -859,7 +859,7 @@ class AsyncClient: async def get( self, url: str, - content: bytes | bytearray | list[int] | None = None, + content: bytes | bytearray | list[int] | Iterable[bytes] | AsyncIterable[bytes] | None = None, data: dict[str, str] | None = None, headers: dict[str, str] | None = None, timeout: float | str | None = USE_CLIENT_DEFAULT, @@ -879,7 +879,7 @@ class AsyncClient: async def post( self, url: str, - content: bytes | bytearray | list[int] | None = None, + content: bytes | bytearray | list[int] | Iterable[bytes] | AsyncIterable[bytes] | None = None, data: dict[str, str] | None = None, headers: dict[str, str] | None = None, timeout: float | str | None = USE_CLIENT_DEFAULT, @@ -900,7 +900,7 @@ class AsyncClient: async def put( self, url: str, - content: bytes | bytearray | list[int] | None = None, + content: bytes | bytearray | list[int] | Iterable[bytes] | AsyncIterable[bytes] | None = None, data: dict[str, str] | None = None, headers: dict[str, str] | None = None, timeout: float | str | None = USE_CLIENT_DEFAULT, @@ -920,7 +920,7 @@ class AsyncClient: async def patch( self, url: str, - content: bytes | bytearray | list[int] | None = None, + content: bytes | bytearray | list[int] | Iterable[bytes] | AsyncIterable[bytes] | None = None, data: dict[str, str] | None = None, headers: dict[str, str] | None = None, timeout: float | str | None = USE_CLIENT_DEFAULT, @@ -940,7 +940,7 @@ class AsyncClient: async def delete( self, url: str, - content: bytes | bytearray | list[int] | None = None, + content: bytes | bytearray | list[int] | Iterable[bytes] | AsyncIterable[bytes] | None = None, data: dict[str, str] | None = None, headers: dict[str, str] | None = None, timeout: float | str | None = USE_CLIENT_DEFAULT, @@ -960,7 +960,7 @@ class AsyncClient: async def head( self, url: str, - content: bytes | bytearray | list[int] | None = None, + content: bytes | bytearray | list[int] | Iterable[bytes] | AsyncIterable[bytes] | None = None, data: dict[str, str] | None = None, headers: dict[str, str] | None = None, timeout: float | str | None = USE_CLIENT_DEFAULT, @@ -980,7 +980,7 @@ class AsyncClient: async def options( self, url: str, - content: bytes | bytearray | list[int] | None = None, + content: bytes | bytearray | list[int] | Iterable[bytes] | AsyncIterable[bytes] | None = None, data: dict[str, str] | None = None, headers: dict[str, str] | None = None, timeout: float | str | None = USE_CLIENT_DEFAULT, @@ -1000,7 +1000,7 @@ class AsyncClient: async def trace( self, url: str, - content: bytes | bytearray | list[int] | None = None, + content: bytes | bytearray | list[int] | Iterable[bytes] | AsyncIterable[bytes] | None = None, data: dict[str, str] | None = None, headers: dict[str, str] | None = None, timeout: float | str | None = USE_CLIENT_DEFAULT, @@ -1021,7 +1021,7 @@ class AsyncClient: self, method: str, url: str, - content: bytes | bytearray | list[int] | None = None, + content: bytes | bytearray | list[int] | Iterable[bytes] | AsyncIterable[bytes] | None = None, data: dict[str, str] | None = None, headers: dict[str, str] | None = None, timeout: float | str | None = USE_CLIENT_DEFAULT, @@ -1045,7 +1045,7 @@ class AsyncClient: self, method: str, url: str, - content: bytes | bytearray | list[int] | None = None, + content: bytes | bytearray | list[int] | Iterable[bytes] | AsyncIterable[bytes] | None = None, data: dict[str, str] | None = None, headers: dict[str, str] | None = None, timeout: float | str | None = USE_CLIENT_DEFAULT, @@ -1078,7 +1078,7 @@ class AsyncClient: def stream( method: str, url: str, - content: bytes | bytearray | list[int] | None = None, + content: bytes | bytearray | list[int] | Iterable[bytes] | None = None, data: dict[str, str] | None = None, headers: dict[str, str] | None = None, timeout: float | str | None = USE_CLIENT_DEFAULT, @@ -1111,7 +1111,7 @@ def stream( def get( url: str, - content: bytes | bytearray | list[int] | None = None, + content: bytes | bytearray | list[int] | Iterable[bytes] | None = None, data: dict[str, str] | None = None, headers: dict[str, str] | None = None, timeout: float | str | None = USE_CLIENT_DEFAULT, @@ -1144,7 +1144,7 @@ def get( def post( url: str, - content: bytes | bytearray | list[int] | None = None, + content: bytes | bytearray | list[int] | Iterable[bytes] | None = None, data: dict[str, str] | None = None, headers: dict[str, str] | None = None, timeout: float | str | None = USE_CLIENT_DEFAULT, @@ -1177,7 +1177,7 @@ def post( def put( url: str, - content: bytes | bytearray | list[int] | None = None, + content: bytes | bytearray | list[int] | Iterable[bytes] | None = None, data: dict[str, str] | None = None, headers: dict[str, str] | None = None, timeout: float | str | None = USE_CLIENT_DEFAULT, @@ -1210,7 +1210,7 @@ def put( def patch( url: str, - content: bytes | bytearray | list[int] | None = None, + content: bytes | bytearray | list[int] | Iterable[bytes] | None = None, data: dict[str, str] | None = None, headers: dict[str, str] | None = None, timeout: float | str | None = USE_CLIENT_DEFAULT, @@ -1243,7 +1243,7 @@ def patch( def delete( url: str, - content: bytes | bytearray | list[int] | None = None, + content: bytes | bytearray | list[int] | Iterable[bytes] | None = None, data: dict[str, str] | None = None, headers: dict[str, str] | None = None, timeout: float | str | None = USE_CLIENT_DEFAULT, @@ -1276,7 +1276,7 @@ def delete( def head( url: str, - content: bytes | bytearray | list[int] | None = None, + content: bytes | bytearray | list[int] | Iterable[bytes] | None = None, data: dict[str, str] | None = None, headers: dict[str, str] | None = None, timeout: float | str | None = USE_CLIENT_DEFAULT, @@ -1309,7 +1309,7 @@ def head( def options( url: str, - content: bytes | bytearray | list[int] | None = None, + content: bytes | bytearray | list[int] | Iterable[bytes] | None = None, data: dict[str, str] | None = None, headers: dict[str, str] | None = None, timeout: float | str | None = USE_CLIENT_DEFAULT, @@ -1339,7 +1339,7 @@ def options( def trace( url: str, - content: bytes | bytearray | list[int] | None = None, + content: bytes | bytearray | list[int] | Iterable[bytes] | None = None, data: dict[str, str] | None = None, headers: dict[str, str] | None = None, timeout: float | str | None = USE_CLIENT_DEFAULT, diff --git a/impit-python/src/async_client.rs b/impit-python/src/async_client.rs index 24f7a6be..646c0903 100644 --- a/impit-python/src/async_client.rs +++ b/impit-python/src/async_client.rs @@ -6,12 +6,12 @@ use impit::{ impit::{Impit, ImpitBuilder}, request::RequestOptions, }; -use pyo3::{exceptions::PyTypeError, ffi::c_str, prelude::*}; +use pyo3::{ffi::c_str, prelude::*}; use crate::{ cookies::PythonCookieJar, errors::ImpitPyError, - request::{form_to_bytes, parse_timeout, RequestBody, USE_CLIENT_DEFAULT_SENTINEL}, + request::{parse_timeout, to_body, RequestBody, USE_CLIENT_DEFAULT_SENTINEL}, response::ImpitPyResponse, }; @@ -135,7 +135,7 @@ impl AsyncClient { &self, py: Python<'python>, url: String, - content: Option>, + content: Option, data: Option, headers: Option>, timeout: Option>, @@ -159,7 +159,7 @@ impl AsyncClient { &self, py: Python<'python>, url: String, - content: Option>, + content: Option, data: Option, headers: Option>, timeout: Option>, @@ -183,7 +183,7 @@ impl AsyncClient { &self, py: Python<'python>, url: String, - content: Option>, + content: Option, data: Option, headers: Option>, timeout: Option>, @@ -207,7 +207,7 @@ impl AsyncClient { &self, py: Python<'python>, url: String, - content: Option>, + content: Option, data: Option, headers: Option>, timeout: Option>, @@ -231,7 +231,7 @@ impl AsyncClient { &self, py: Python<'python>, url: String, - content: Option>, + content: Option, data: Option, headers: Option>, timeout: Option>, @@ -255,7 +255,7 @@ impl AsyncClient { &self, py: Python<'python>, url: String, - content: Option>, + content: Option, data: Option, headers: Option>, timeout: Option>, @@ -279,7 +279,7 @@ impl AsyncClient { &self, py: Python<'python>, url: String, - content: Option>, + content: Option, data: Option, headers: Option>, timeout: Option>, @@ -303,7 +303,7 @@ impl AsyncClient { &self, py: Python<'python>, url: String, - content: Option>, + content: Option, data: Option, headers: Option>, timeout: Option>, @@ -328,7 +328,7 @@ impl AsyncClient { py: Python<'python>, method: &str, url: String, - content: Option>, + content: Option, data: Option, headers: Option>, timeout: Option>, @@ -374,7 +374,7 @@ impl AsyncClient { py: Python<'python>, method: &str, url: String, - content: Option>, + content: Option, mut data: Option, headers: Option>, timeout: Option>, @@ -384,25 +384,10 @@ impl AsyncClient { let mut headers = headers.clone(); if let Some(content) = content { - data = Some(RequestBody::Bytes(content)); + data = Some(content); } - let body: Vec = match data { - Some(data) => match data { - RequestBody::Bytes(bytes) => Ok(bytes), - RequestBody::Form(form) => { - headers.get_or_insert_with(HashMap::new).insert( - "Content-Type".to_string(), - "application/x-www-form-urlencoded".to_string(), - ); - Ok(form_to_bytes(form)) - } - RequestBody::CatchAll(e) => Err(PyErr::new::(format!( - "Unsupported data type in request body: {e:#?}" - ))), - }, - None => Ok(Vec::new()), - }?; + let body = to_body(data, &mut headers)?; let timeout = parse_timeout(timeout)?; @@ -423,14 +408,14 @@ impl AsyncClient { pyo3_async_runtimes::tokio::future_into_py::<_, ImpitPyResponse>(py, async move { let response = match method_str.to_lowercase().as_str() { - "get" => impit.get(url, Some(body.into()), Some(options)).await, - "post" => impit.post(url, Some(body.into()), Some(options)).await, - "patch" => impit.patch(url, Some(body.into()), Some(options)).await, - "put" => impit.put(url, Some(body.into()), Some(options)).await, - "options" => impit.options(url, Some(body.into()), Some(options)).await, - "trace" => impit.trace(url, Some(body.into()), Some(options)).await, - "head" => impit.head(url, Some(body.into()), Some(options)).await, - "delete" => impit.delete(url, Some(body.into()), Some(options)).await, + "get" => impit.get(url, Some(body), Some(options)).await, + "post" => impit.post(url, Some(body), Some(options)).await, + "patch" => impit.patch(url, Some(body), Some(options)).await, + "put" => impit.put(url, Some(body), Some(options)).await, + "options" => impit.options(url, Some(body), Some(options)).await, + "trace" => impit.trace(url, Some(body), Some(options)).await, + "head" => impit.head(url, Some(body), Some(options)).await, + "delete" => impit.delete(url, Some(body), Some(options)).await, _ => Err(ImpitError::InvalidMethod(method_str.to_string())), }; diff --git a/impit-python/src/client.rs b/impit-python/src/client.rs index 06db785f..e691f93b 100644 --- a/impit-python/src/client.rs +++ b/impit-python/src/client.rs @@ -11,7 +11,7 @@ use pyo3::{ffi::c_str, prelude::*}; use crate::{ cookies::PythonCookieJar, errors::ImpitPyError, - request::{form_to_bytes, parse_timeout, RequestBody, USE_CLIENT_DEFAULT_SENTINEL}, + request::{parse_timeout, to_body, RequestBody, USE_CLIENT_DEFAULT_SENTINEL}, response::{self, ImpitPyResponse}, }; @@ -132,7 +132,7 @@ impl Client { &self, py: Python<'_>, url: String, - content: Option>, + content: Option, data: Option, headers: Option>, timeout: Option>, @@ -156,7 +156,7 @@ impl Client { &self, py: Python<'_>, url: String, - content: Option>, + content: Option, data: Option, headers: Option>, timeout: Option>, @@ -180,7 +180,7 @@ impl Client { &self, py: Python<'_>, url: String, - content: Option>, + content: Option, data: Option, headers: Option>, timeout: Option>, @@ -204,7 +204,7 @@ impl Client { &self, py: Python<'_>, url: String, - content: Option>, + content: Option, data: Option, headers: Option>, timeout: Option>, @@ -228,7 +228,7 @@ impl Client { &self, py: Python<'_>, url: String, - content: Option>, + content: Option, data: Option, headers: Option>, timeout: Option>, @@ -252,7 +252,7 @@ impl Client { &self, py: Python<'_>, url: String, - content: Option>, + content: Option, data: Option, headers: Option>, timeout: Option>, @@ -276,7 +276,7 @@ impl Client { &self, py: Python<'_>, url: String, - content: Option>, + content: Option, data: Option, headers: Option>, timeout: Option>, @@ -300,7 +300,7 @@ impl Client { &self, py: Python<'_>, url: String, - content: Option>, + content: Option, data: Option, headers: Option>, timeout: Option>, @@ -325,7 +325,7 @@ impl Client { py: Python<'python>, method: &str, url: String, - content: Option>, + content: Option, data: Option, headers: Option>, timeout: Option>, @@ -371,7 +371,7 @@ impl Client { py: Python<'_>, method: &str, url: String, - content: Option>, + content: Option, mut data: Option, headers: Option>, timeout: Option>, @@ -381,25 +381,11 @@ impl Client { let mut headers = headers.clone(); if let Some(content) = content { - data = Some(RequestBody::Bytes(content)); + data = Some(content); } - let body: Vec = match data { - Some(data) => match data { - RequestBody::Bytes(bytes) => Ok(bytes), - RequestBody::Form(form) => { - headers.get_or_insert_with(HashMap::new).insert( - "Content-Type".to_string(), - "application/x-www-form-urlencoded".to_string(), - ); - Ok(form_to_bytes(form)) - } - RequestBody::CatchAll(e) => Err(ImpitPyError(ImpitError::BindingPassthroughError( - format!("Unsupported data type: {e:?}").to_string(), - ))), - }, - None => Ok(Vec::new()), - }?; + let body = to_body(data, &mut headers) + .map_err(|e| ImpitPyError(ImpitError::BindingPassthroughError(e.to_string())))?; let timeout = parse_timeout(timeout) .map_err(|e| ImpitPyError(ImpitError::BindingPassthroughError(e.to_string())))?; @@ -417,30 +403,14 @@ impl Client { py.detach(|| { pyo3_async_runtimes::tokio::get_runtime().block_on(async { let response = match method.to_lowercase().as_str() { - "get" => self.impit.get(url, Some(body.into()), Some(options)).await, - "post" => self.impit.post(url, Some(body.into()), Some(options)).await, - "patch" => { - self.impit - .patch(url, Some(body.into()), Some(options)) - .await - } - "put" => self.impit.put(url, Some(body.into()), Some(options)).await, - "options" => { - self.impit - .options(url, Some(body.into()), Some(options)) - .await - } - "trace" => { - self.impit - .trace(url, Some(body.into()), Some(options)) - .await - } - "head" => self.impit.head(url, Some(body.into()), Some(options)).await, - "delete" => { - self.impit - .delete(url, Some(body.into()), Some(options)) - .await - } + "get" => self.impit.get(url, Some(body), Some(options)).await, + "post" => self.impit.post(url, Some(body), Some(options)).await, + "patch" => self.impit.patch(url, Some(body), Some(options)).await, + "put" => self.impit.put(url, Some(body), Some(options)).await, + "options" => self.impit.options(url, Some(body), Some(options)).await, + "trace" => self.impit.trace(url, Some(body), Some(options)).await, + "head" => self.impit.head(url, Some(body), Some(options)).await, + "delete" => self.impit.delete(url, Some(body), Some(options)).await, _ => Err(ImpitError::InvalidMethod(method.to_string())), }; diff --git a/impit-python/src/lib.rs b/impit-python/src/lib.rs index 9e16af38..e6b89936 100644 --- a/impit-python/src/lib.rs +++ b/impit-python/src/lib.rs @@ -93,7 +93,7 @@ fn impit(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { fn $name( _py: Python, url: String, - content: Option>, + content: Option, data: Option, headers: Option>, timeout: Option>, @@ -122,7 +122,7 @@ fn impit(_py: Python, m: &Bound<'_, PyModule>) -> PyResult<()> { _py: Python<'python>, method: &str, url: String, - content: Option>, + content: Option, data: Option, headers: Option>, timeout: Option>, diff --git a/impit-python/src/request.rs b/impit-python/src/request.rs index 07a26da6..4e9102dc 100644 --- a/impit-python/src/request.rs +++ b/impit-python/src/request.rs @@ -1,7 +1,15 @@ -use std::{collections::HashMap, time::Duration}; +use std::{collections::HashMap, io, time::Duration}; +use bytes::Bytes; use either::{Either, Left, Right}; -use pyo3::{Bound, PyAny}; +use futures::{stream, Stream}; +use impit::request::ImpitBody; +use pyo3::{ + exceptions::{PyStopAsyncIteration, PyStopIteration, PyTypeError}, + types::PyAnyMethods, + Borrowed, Bound, Py, PyAny, PyErr, PyResult, PyTypeInfo, Python, +}; +use pyo3_async_runtimes::TaskLocals; /// The sentinel string used as the Python default value for per-request `timeout` parameters. /// @@ -36,13 +44,109 @@ pub(crate) fn parse_timeout( use pyo3::FromPyObject; #[derive(FromPyObject)] -pub(crate) enum RequestBody<'py> { +pub(crate) enum RequestBody { #[pyo3(transparent, annotation = "bytes")] Bytes(Vec), #[pyo3(transparent, annotation = "dict[str, str]")] Form(HashMap), + #[pyo3(transparent, annotation = "Iterable[bytes] | AsyncIterable[bytes]")] + Iterator(PyIterator), #[pyo3(transparent)] - CatchAll(Bound<'py, PyAny>), // This extraction never fails + CatchAll(Py), // This extraction never fails +} + +/// A Python iterator over the chunks of a request body. +pub(crate) enum PyIterator { + Sync(Py), + Async(Py, TaskLocals), +} + +impl<'py> FromPyObject<'_, 'py> for PyIterator { + type Error = PyErr; + + fn extract(object: Borrowed<'_, 'py, PyAny>) -> PyResult { + match object.call_method0("__aiter__") { + Ok(iterator) => Ok(Self::Async( + iterator.unbind(), + TaskLocals::with_running_loop(object.py())?.copy_context(object.py())?, + )), + Err(_) => Ok(Self::Sync(object.try_iter()?.into_any().unbind())), + } + } +} + +fn to_chunk(next: PyResult>) -> Option> { + Python::attach(|py| match next { + Ok(chunk) => Some( + chunk + .extract::>(py) + .map(Bytes::from) + .map_err(io::Error::other), + ), + Err(err) if err.is_instance_of::(py) => None, + Err(err) => Some(Err(io::Error::other(err))), + }) +} + +fn to_stream(iterator: PyIterator) -> impl Stream> { + stream::unfold(iterator, |iterator| async move { + let chunk = match &iterator { + PyIterator::Sync(iterator) => { + let iterator = Python::attach(|py| iterator.clone_ref(py)); + tokio::task::spawn_blocking(move || { + to_chunk::(Python::attach(|py| { + iterator + .bind(py) + .call_method0("__next__") + .map(Bound::unbind) + })) + }) + .await + .unwrap_or_else(|err| Some(Err(io::Error::other(err)))) + } + PyIterator::Async(iterator, locals) => { + let next = Python::attach(|py| { + pyo3_async_runtimes::into_future_with_locals( + locals, + iterator.bind(py).call_method0("__anext__")?, + ) + }); + to_chunk::(match next { + Ok(next) => next.await, + Err(err) => Err(err), + }) + } + }; + + Some((chunk?, iterator)) + }) +} + +/// Converts the Python request body into an [`ImpitBody`], streaming it if it is an iterator. +pub(crate) fn to_body( + data: Option, + headers: &mut Option>, +) -> PyResult { + Ok(match data { + None => ImpitBody::Empty, + Some(RequestBody::Bytes(bytes)) => bytes.into(), + Some(RequestBody::Form(form)) => { + headers.get_or_insert_default().insert( + "Content-Type".to_string(), + "application/x-www-form-urlencoded".to_string(), + ); + form_to_bytes(form).into() + } + Some(RequestBody::Iterator(iterator)) => ImpitBody::from_stream(to_stream(iterator)), + Some(RequestBody::CatchAll(object)) => { + return Err(Python::attach(|py| { + PyErr::new::(format!( + "Unsupported data type in request body: {}", + object.bind(py).get_type() + )) + })) + } + }) } pub fn form_to_bytes(data: HashMap) -> Vec { diff --git a/impit-python/test/async_client_test.py b/impit-python/test/async_client_test.py index b062e305..7bcf6ea3 100644 --- a/impit-python/test/async_client_test.py +++ b/impit-python/test/async_client_test.py @@ -2,6 +2,7 @@ import json import socket import threading +from collections.abc import AsyncIterator from http.cookiejar import Cookie, CookieJar from typing import Literal @@ -526,6 +527,19 @@ async def test_passing_string_body(self, browser: Browser) -> None: assert response.status_code == 200 assert json.loads(response.text)['data'] == '{"Impit-Test":"foořžš"}' + @pytest.mark.asyncio + async def test_passing_async_iterator_body(self, browser: Browser) -> None: + impit = AsyncClient(browser=browser) + + async def chunks() -> AsyncIterator[bytes]: + yield b'{"Impit-Test":' + yield b'"foo"}' + + response = await impit.post(get_httpbin_url('/post'), content=chunks()) + assert response.status_code == 200 + assert json.loads(response.text)['data'] == '{"Impit-Test":"foo"}' + assert json.loads(response.text)['headers']['Transfer-Encoding'] == 'chunked' + @pytest.mark.asyncio async def test_passing_string_body_in_data(self, browser: Browser) -> None: impit = AsyncClient(browser=browser) diff --git a/impit-python/test/basic_client_test.py b/impit-python/test/basic_client_test.py index f78bffe6..f1e1f80a 100644 --- a/impit-python/test/basic_client_test.py +++ b/impit-python/test/basic_client_test.py @@ -2,6 +2,7 @@ import socket import threading import time +from collections.abc import Iterator from http.cookiejar import Cookie, CookieJar from typing import Literal @@ -454,6 +455,18 @@ def test_passing_string_body(self, browser: Browser) -> None: assert response.status_code == 200 assert response.json()['data'] == '{"Impit-Test":"foořžš"}' + def test_passing_iterator_body(self, browser: Browser) -> None: + impit = Client(browser=browser) + + def chunks() -> Iterator[bytes]: + yield b'{"Impit-Test":' + yield b'"foo"}' + + response = impit.post(get_httpbin_url('/post'), content=chunks()) + assert response.status_code == 200 + assert response.json()['data'] == '{"Impit-Test":"foo"}' + assert response.json()['headers']['Transfer-Encoding'] == 'chunked' + def test_passing_string_body_in_data(self, browser: Browser) -> None: impit = Client(browser=browser)