From de17fef60f1db0ed1667fab59048be2b5c34fb8a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jind=C5=99ich=20B=C3=A4r?= Date: Thu, 13 Aug 2026 12:04:50 +0200 Subject: [PATCH] feat(python): stream iterator and async iterator request bodies `content=` now accepts any (async) iterable of byte chunks, which is fed into the request as the chunks are produced instead of being buffered first. Related: #513 --- impit-python/python/impit/impit.pyi | 60 ++++++------- impit-python/src/async_client.rs | 59 +++++-------- impit-python/src/client.rs | 74 +++++----------- impit-python/src/lib.rs | 4 +- impit-python/src/request.rs | 112 ++++++++++++++++++++++++- impit-python/test/async_client_test.py | 14 ++++ impit-python/test/basic_client_test.py | 13 +++ 7 files changed, 211 insertions(+), 125 deletions(-) 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)