From 01c52707c89d05ffd89be3ecd6ee9930dbb36059 Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Thu, 27 Mar 2025 10:43:54 +0800 Subject: [PATCH 1/2] refactor(typing): refactor type extractors --- examples/impersonate.py | 82 +++++++++++++++++++----- rnet.pyi | 24 ++++--- src/async_impl/client.rs | 44 +++---------- src/async_impl/request.rs | 4 +- src/blocking/client.rs | 10 +-- src/lib.rs | 5 +- src/typing/body.rs | 14 ++--- src/typing/cookie.rs | 7 +-- src/typing/headers.rs | 12 ++-- src/typing/mod.rs | 118 ++++++++++++++++++++++++++++++++--- src/typing/multipart/form.rs | 14 +++++ src/typing/multipart/mod.rs | 5 +- src/typing/multipart/part.rs | 22 ++++--- src/typing/param/client.rs | 82 +++++------------------- src/typing/param/mod.rs | 9 +++ src/typing/param/request.rs | 41 +++++------- src/typing/param/ws.rs | 19 +++--- src/typing/proxy.rs | 41 ++++++++++-- 18 files changed, 339 insertions(+), 214 deletions(-) diff --git a/examples/impersonate.py b/examples/impersonate.py index 4b661461..ae2160ca 100644 --- a/examples/impersonate.py +++ b/examples/impersonate.py @@ -1,28 +1,76 @@ import asyncio -from rnet import Impersonate, Client +from rnet import Impersonate, ImpersonateOS, ImpersonateOption, Client, Response -async def main(): - headers = {"foo": "bar", "bar": "foo"} - headers_order = ["accept-encoding", "foo", "bar"] +async def print_response_info(resp: Response): + """Helper function to print response details + + Args: + resp: Response object from the request + """ + async with resp: + print("\n=== Response Information ===") + print(f"Status Code: {resp.status_code}") + print(f"Version: {resp.version}") + print(f"Response URL: {resp.url}") + print(f"Headers: {resp.headers}") + print(f"Encoding: {resp.encoding}") + print(f"Content-Length: {resp.content_length}") + print(f"Remote Address: {resp.remote_addr}") + print(f"Peer Certificate: {resp.peer_certificate()}") + print(f"Content: {await resp.text()}") + print("========================\n") + + +async def request_firefox(): + """Test request using Firefox browser impersonation + + Demonstrates basic browser impersonation with custom header order + """ + print("\n[Testing Firefox Impersonation]") client = Client( impersonate=Impersonate.Firefox135, - user_agent="rnet", + headers_order=["accept-encoding", "user-agent", "accept"], tls_info=True, - default_headers=headers, - headers_order=headers_order, ) - async with await client.get("https://tls.peet.ws/api/all") as resp: - print("Status Code: ", resp.status_code) - print("Version: ", resp.version) - print("Response URL: ", resp.url) - print("Headers: ", resp.headers) - print("Encoding: ", resp.encoding) - print("Content-Length: ", resp.content_length) - print("Remote Address: ", resp.remote_addr) - print("Peer Certificate: ", resp.peer_certificate()) - print("Content: ", await resp.text()) + resp = await client.get("https://tls.peet.ws/api/all") + await print_response_info(resp) + return client + + +async def request_chrome_android(client: Client): + """Test request using Chrome on Android impersonation + + Demonstrates advanced impersonation with OS specification + + Args: + client: Existing client instance to update + """ + print("\n[Testing Chrome on Android Impersonation]") + client.update( + impersonate=ImpersonateOption( + impersonate=Impersonate.Chrome134, + impersonate_os=ImpersonateOS.Android, + ) + ) + resp = await client.get("https://tls.peet.ws/api/all") + await print_response_info(resp) + + +async def main(): + """Main function to run the impersonation examples + + Demonstrates different browser impersonation scenarios: + 1. Firefox with custom header order + 2. Chrome on Android with OS specification + """ + # First test with Firefox + client = await request_firefox() + + # Then update and test with Chrome on Android + await request_chrome_android(client) if __name__ == "__main__": + # Run the async main function asyncio.run(main()) diff --git a/rnet.pyi b/rnet.pyi index 045274e6..e315ca1a 100644 --- a/rnet.pyi +++ b/rnet.pyi @@ -78,10 +78,7 @@ class BlockingClient: # Arguments * `**kwds` - The parameters to update the client with. - impersonate: typing.Optional[Impersonate] - impersonate_os: typing.Optional[ImpersonateOS] - impersonate_skip_http2: typing.Optional[builtins.bool] - impersonate_skip_headers: typing.Optional[builtins.bool] + impersonate: typing.Optional[typing.Union[Impersonate, ImpersonateOption]] headers: typing.Optional[typing.Dict[str, bytes]] headers_order: typing.Optional[typing.List[str]] proxies: typing.Optional[builtins.list[Proxy]] @@ -858,10 +855,7 @@ class Client: # Arguments * `**kwds` - The parameters to update the client with. - impersonate: typing.Optional[Impersonate] - impersonate_os: typing.Optional[ImpersonateOS] - impersonate_skip_http2: typing.Optional[builtins.bool] - impersonate_skip_headers: typing.Optional[builtins.bool] + impersonate: typing.Optional[typing.Union[Impersonate, ImpersonateOption]] headers: typing.Optional[typing.Dict[str, bytes]] headers_order: typing.Optional[typing.List[str]] proxies: typing.Optional[builtins.list[Proxy]] @@ -1439,6 +1433,20 @@ class HeaderMapKeysIter: def __iter__(self) -> HeaderMapKeysIter: ... def __next__(self) -> typing.Optional[typing.Any]: ... +class ImpersonateOption: + r""" + A struct to represent the `ImpersonateOption` class. + """ + + def __new__( + cls, + impersonate: Impersonate, + impersonate_os: typing.Optional[ImpersonateOS] = None, + skip_http2: typing.Optional[builtins.bool] = None, + skip_headers: typing.Optional[builtins.bool] = None, + ): ... + ... + class Message: r""" A WebSocket message. diff --git a/src/async_impl/client.rs b/src/async_impl/client.rs index 55319d82..a4afda25 100644 --- a/src/async_impl/client.rs +++ b/src/async_impl/client.rs @@ -5,7 +5,7 @@ use crate::{ dns, error::Error, typing::{ - Cookie, HeaderMap, ImpersonateOS, Method, SslVerify, TlsVersion, + Cookie, HeaderMap, Method, SslVerify, TlsVersion, param::{ClientParams, RequestParams, UpdateClientParams, WebSocketParams}, }, }; @@ -361,10 +361,7 @@ impl Client { /// /// * `**kwds` - Optional request parameters as a dictionary. /// - /// impersonate: typing.Optional[Impersonate] - /// impersonate_os: typing.Optional[ImpersonateOS] - /// impersonate_skip_http2: typing.Optional[builtins.bool] - /// impersonate_skip_headers: typing.Optional[builtins.bool] + /// impersonate: typing.Optional[typing.Union[Impersonate, ImpersonateOption]] /// base_url: typing.Optional[str] /// user_agent: typing.Optional[str] /// default_headers: typing.Optional[typing.Dict[str, bytes]] @@ -426,19 +423,7 @@ impl Client { // Impersonation options. if let Some(impersonate) = params.impersonate.take() { - builder = builder.emulation( - rquest_util::EmulationOption::builder() - .emulation(impersonate.into_ffi()) - .emulation_os( - params - .impersonate_os - .map(ImpersonateOS::into_ffi) - .unwrap_or_default(), - ) - .skip_http2(params.impersonate_skip_http2.unwrap_or(false)) - .skip_headers(params.impersonate_skip_headers.unwrap_or(false)) - .build(), - ); + builder = builder.emulation(impersonate.0); } // User agent options. @@ -600,7 +585,7 @@ impl Client { // Network options. if let Some(proxies) = params.proxies.take() { - for proxy in proxies { + for proxy in proxies.0 { builder = builder.proxy(proxy); } } @@ -779,10 +764,7 @@ impl Client { /// # Arguments /// * `**kwds` - The parameters to update the client with. /// - /// impersonate: typing.Optional[Impersonate] - /// impersonate_os: typing.Optional[ImpersonateOS] - /// impersonate_skip_http2: typing.Optional[builtins.bool] - /// impersonate_skip_headers: typing.Optional[builtins.bool] + /// impersonate: typing.Optional[typing.Union[Impersonate, ImpersonateOption]] /// headers: typing.Optional[typing.Dict[str, bytes]] /// headers_order: typing.Optional[typing.List[str]] /// proxies: typing.Optional[builtins.list[Proxy]] @@ -809,19 +791,7 @@ impl Client { // Impersonation options. if let Some(impersonate) = params.impersonate.take() { - update = update.emulation( - rquest_util::EmulationOption::builder() - .emulation(impersonate.into_ffi()) - .emulation_os( - params - .impersonate_os - .map(ImpersonateOS::into_ffi) - .unwrap_or_default(), - ) - .skip_http2(params.impersonate_skip_http2.unwrap_or(false)) - .skip_headers(params.impersonate_skip_headers.unwrap_or(false)) - .build(), - ); + update = update.emulation(impersonate.0); } // Default headers options. @@ -838,7 +808,7 @@ impl Client { ); // Network options. - apply_option!(apply_if_some, update, params.proxies, proxies); + apply_option!(apply_if_some_inner, update, params.proxies, proxies); apply_option!( apply_transformed_option, update, diff --git a/src/async_impl/request.rs b/src/async_impl/request.rs index f6b0b46d..a224e0b1 100644 --- a/src/async_impl/request.rs +++ b/src/async_impl/request.rs @@ -65,7 +65,7 @@ where ); // Network options. - apply_option!(apply_if_some, builder, params.proxy, proxy); + apply_option!(apply_if_some_inner, builder, params.proxy, proxy); apply_option!( apply_transformed_option, builder, @@ -123,7 +123,7 @@ where apply_option!(apply_if_some, builder, params.body, body); // Multipart options. - apply_option!(apply_if_some, builder, params.multipart, multipart); + apply_option!(apply_if_some_inner, builder, params.multipart, multipart); // Send the request. builder diff --git a/src/blocking/client.rs b/src/blocking/client.rs index 1edbbb10..b50ff2fb 100644 --- a/src/blocking/client.rs +++ b/src/blocking/client.rs @@ -351,10 +351,7 @@ impl BlockingClient { /// /// * `**kwds` - Optional request parameters as a dictionary. /// - /// impersonate: typing.Optional[Impersonate] - /// impersonate_os: typing.Optional[ImpersonateOS] - /// impersonate_skip_http2: typing.Optional[builtins.bool] - /// impersonate_skip_headers: typing.Optional[builtins.bool] + /// impersonate: typing.Optional[typing.Union[Impersonate, ImpersonateOption]] /// user_agent: typing.Optional[str] /// default_headers: typing.Optional[typing.Dict[str, bytes]] /// headers_order: typing.Optional[typing.List[str]] @@ -504,10 +501,7 @@ impl BlockingClient { /// # Arguments /// * `**kwds` - The parameters to update the client with. /// - /// impersonate: typing.Optional[Impersonate] - /// impersonate_os: typing.Optional[ImpersonateOS] - /// impersonate_skip_http2: typing.Optional[builtins.bool] - /// impersonate_skip_headers: typing.Optional[builtins.bool] + /// impersonate: typing.Optional[typing.Union[Impersonate, ImpersonateOption]] /// headers: typing.Optional[typing.Dict[str, bytes]] /// headers_order: typing.Optional[typing.List[str]] /// proxies: typing.Optional[builtins.list[Proxy]] diff --git a/src/lib.rs b/src/lib.rs index 5ea5e832..5d7b121f 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -15,8 +15,8 @@ use pyo3_stub_gen::{define_stub_info_gatherer, derive::*}; use typing::param::{RequestParams, WebSocketParams}; use typing::{ Cookie, HeaderMap, HeaderMapItemsIter, HeaderMapKeysIter, Impersonate, ImpersonateOS, - LookupIpStrategy, Method, Multipart, Part, Proxy, SameSite, SocketAddr, StatusCode, TlsVersion, - Version, + ImpersonateOption, LookupIpStrategy, Method, Multipart, Part, Proxy, SameSite, SocketAddr, + StatusCode, TlsVersion, Version, }; #[cfg(not(target_env = "msvc"))] @@ -325,6 +325,7 @@ fn rnet(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/src/typing/body.rs b/src/typing/body.rs index 61855577..7f010152 100644 --- a/src/typing/body.rs +++ b/src/typing/body.rs @@ -6,24 +6,24 @@ use pyo3::{FromPyObject, PyAny}; use rquest::Body; /// The body to use for the request. -pub enum FromPyBody { +pub enum BodyExtractor { Text(Bytes), Bytes(Bytes), SyncStream(SyncStream), AsyncStream(AsyncStream), } -impl From for Body { - fn from(value: FromPyBody) -> Body { +impl From for Body { + fn from(value: BodyExtractor) -> Body { match value { - FromPyBody::Text(bytes) | FromPyBody::Bytes(bytes) => Body::from(bytes), - FromPyBody::SyncStream(stream) => Body::wrap_stream(stream), - FromPyBody::AsyncStream(stream) => Body::wrap_stream(stream), + BodyExtractor::Text(bytes) | BodyExtractor::Bytes(bytes) => Body::from(bytes), + BodyExtractor::SyncStream(stream) => Body::wrap_stream(stream), + BodyExtractor::AsyncStream(stream) => Body::wrap_stream(stream), } } } -impl FromPyObject<'_> for FromPyBody { +impl FromPyObject<'_> for BodyExtractor { fn extract_bound(ob: &Bound<'_, PyAny>) -> PyResult { if let Ok(text) = ob.extract::() { return Ok(Self::Text(Bytes::from_owner(text))); diff --git a/src/typing/cookie.rs b/src/typing/cookie.rs index ce421599..6e970f0b 100644 --- a/src/typing/cookie.rs +++ b/src/typing/cookie.rs @@ -173,10 +173,9 @@ impl Cookie { } } -/// Parse a cookie header from a Python dictionary. -pub struct CookieFromPyDict(pub HeaderValue); +pub struct CookieExtractor(pub HeaderValue); -impl FromPyObject<'_> for CookieFromPyDict { +impl FromPyObject<'_> for CookieExtractor { fn extract_bound(ob: &Bound<'_, PyAny>) -> PyResult { let dict = ob.downcast::()?; dict.iter() @@ -202,7 +201,7 @@ impl FromPyObject<'_> for CookieFromPyDict { } #[cfg(feature = "docs")] -impl PyStubType for CookieFromPyDict { +impl PyStubType for CookieExtractor { fn type_output() -> TypeInfo { TypeInfo::with_module("typing.Dict[str, str]", "typing".into()) } diff --git a/src/typing/headers.rs b/src/typing/headers.rs index e1babb75..fc017836 100644 --- a/src/typing/headers.rs +++ b/src/typing/headers.rs @@ -159,12 +159,12 @@ impl HeaderMapItemsIter { } /// A HTTP header map. -pub struct HeaderMapFromPy(pub header::HeaderMap); +pub struct HeaderMapExtractor(pub header::HeaderMap); /// A list of header names in order. -pub struct HeadersOrderFromPyList(pub Vec); +pub struct HeadersOrderExtractor(pub Vec); -impl FromPyObject<'_> for HeaderMapFromPy { +impl FromPyObject<'_> for HeaderMapExtractor { fn extract_bound(ob: &Bound<'_, PyAny>) -> PyResult { if let Ok(headers) = ob.downcast::() { return Ok(Self(headers.borrow().0.clone())); @@ -188,7 +188,7 @@ impl FromPyObject<'_> for HeaderMapFromPy { } #[cfg(feature = "docs")] -impl<'py> IntoPyObject<'py> for HeaderMapFromPy { +impl<'py> IntoPyObject<'py> for HeaderMapExtractor { type Target = HeaderMap; type Output = Bound<'py, Self::Target>; @@ -201,7 +201,7 @@ impl<'py> IntoPyObject<'py> for HeaderMapFromPy { } #[cfg(feature = "docs")] -impl PyStubType for HeaderMapFromPy { +impl PyStubType for HeaderMapExtractor { fn type_output() -> TypeInfo { TypeInfo::with_module( "typing.Union[typing.Dict[str, str], HeaderMap]", @@ -210,7 +210,7 @@ impl PyStubType for HeaderMapFromPy { } } -impl<'py> FromPyObject<'py> for HeadersOrderFromPyList { +impl<'py> FromPyObject<'py> for HeadersOrderExtractor { fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { let list = ob.downcast::()?; list.iter() diff --git a/src/typing/mod.rs b/src/typing/mod.rs index c9afa865..12532bf4 100644 --- a/src/typing/mod.rs +++ b/src/typing/mod.rs @@ -11,25 +11,28 @@ mod ssl; mod status; pub use self::{ - body::FromPyBody, - cookie::{Cookie, CookieFromPyDict}, + body::BodyExtractor, + cookie::{Cookie, CookieExtractor}, enums::{Impersonate, ImpersonateOS, LookupIpStrategy, Method, SameSite, TlsVersion, Version}, headers::{ - HeaderMap, HeaderMapFromPy, HeaderMapItemsIter, HeaderMapKeysIter, HeadersOrderFromPyList, + HeaderMap, HeaderMapExtractor, HeaderMapItemsIter, HeaderMapKeysIter, HeadersOrderExtractor, }, ipaddr::{IpAddr, SocketAddr}, json::Json, multipart::{Multipart, Part}, - proxy::Proxy, + proxy::{Proxy, ProxyExtractor}, ssl::SslVerify, status::StatusCode, }; use pyo3::{prelude::*, pybacked::PyBackedStr}; +#[cfg(feature = "docs")] +use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods}; +use rquest_util::EmulationOption; use serde::ser::{Serialize, SerializeSeq, Serializer}; -pub struct QueryOrForm(Vec<(PyBackedStr, PyBackedStr)>); +pub struct UrlEncodedValuesExtractor(Vec<(PyBackedStr, PyBackedStr)>); -impl Serialize for QueryOrForm { +impl Serialize for UrlEncodedValuesExtractor { fn serialize(&self, serializer: S) -> Result where S: Serializer, @@ -42,8 +45,109 @@ impl Serialize for QueryOrForm { } } -impl FromPyObject<'_> for QueryOrForm { +impl FromPyObject<'_> for UrlEncodedValuesExtractor { fn extract_bound(ob: &Bound<'_, PyAny>) -> PyResult { ob.extract().map(Self) } } + +pub struct ImpersonateExtractor(pub EmulationOption); + +impl FromPyObject<'_> for ImpersonateExtractor { + fn extract_bound(ob: &Bound<'_, PyAny>) -> PyResult { + if let Ok(impersonate) = ob.downcast::() { + let emulation = EmulationOption::builder() + .emulation(impersonate.borrow().clone().into_ffi()) + .build(); + + return Ok(Self(emulation)); + } + + let option = ob.downcast::()?.borrow(); + + Ok(Self( + EmulationOption::builder() + .emulation(option.impersonate.into_ffi()) + .emulation_os( + option + .impersonate_os + .map(|os| os.into_ffi()) + .unwrap_or_default(), + ) + .skip_http2(option.skip_http2.unwrap_or(false)) + .skip_headers(option.skip_headers.unwrap_or(false)) + .build(), + )) + } +} + +/// A struct to represent the `ImpersonateOption` class. +#[cfg_attr(feature = "docs", gen_stub_pyclass)] +#[pyclass] +pub struct ImpersonateOption { + /// The browser version to impersonate. + impersonate: Impersonate, + + /// The operating system. + impersonate_os: Option, + + /// Whether to skip HTTP/2. + skip_http2: Option, + + /// Whether to skip headers. + skip_headers: Option, +} + +#[cfg_attr(feature = "docs", gen_stub_pymethods)] +#[pymethods] +impl ImpersonateOption { + /// Create a new impersonation option instance. + /// + /// This class allows you to configure browser/client impersonation settings + /// including the browser type, operating system, and HTTP protocol options. + /// + /// Args: + /// impersonate (Impersonate): The browser/client type to impersonate + /// impersonate_os (Optional[ImpersonateOS]): The operating system to impersonate, defaults to None + /// skip_http2 (Optional[bool]): Whether to disable HTTP/2 support, defaults to False + /// skip_headers (Optional[bool]): Whether to skip default request headers, defaults to False + /// + /// Returns: + /// ImpersonateOption: A new impersonation option instance + /// + /// Examples: + /// ```python + /// from rnet import ImpersonateOption, Impersonate, ImpersonateOS + /// + /// # Basic Chrome 120 impersonation + /// option = ImpersonateOption(Impersonate.Chrome120) + /// + /// # Firefox 136 on Windows with custom options + /// option = ImpersonateOption( + /// impersonate=Impersonate.Firefox136, + /// impersonate_os=ImpersonateOS.Windows, + /// skip_http2=False, + /// skip_headers=True + /// ) + /// ``` + #[new] + #[pyo3(signature = ( + impersonate, + impersonate_os = None, + skip_http2 = None, + skip_headers = None + ))] + fn new( + impersonate: Impersonate, + impersonate_os: Option, + skip_http2: Option, + skip_headers: Option, + ) -> Self { + Self { + impersonate, + impersonate_os, + skip_http2, + skip_headers, + } + } +} diff --git a/src/typing/multipart/form.rs b/src/typing/multipart/form.rs index aa92b898..79f6e4d6 100644 --- a/src/typing/multipart/form.rs +++ b/src/typing/multipart/form.rs @@ -31,3 +31,17 @@ impl Multipart { Ok(Multipart(Some(new_form))) } } + +pub struct MultipartExtractor(pub Form); + +impl FromPyObject<'_> for MultipartExtractor { + fn extract_bound(ob: &Bound<'_, PyAny>) -> PyResult { + let form = ob.downcast::()?; + form.borrow_mut() + .0 + .take() + .map(Self) + .ok_or_else(|| Error::MemoryError) + .map_err(Into::into) + } +} diff --git a/src/typing/multipart/mod.rs b/src/typing/multipart/mod.rs index 77b0013f..eea371fe 100644 --- a/src/typing/multipart/mod.rs +++ b/src/typing/multipart/mod.rs @@ -1,4 +1,7 @@ mod form; mod part; -pub use self::{form::Multipart, part::Part}; +pub use self::{ + form::{Multipart, MultipartExtractor}, + part::Part, +}; diff --git a/src/typing/multipart/part.rs b/src/typing/multipart/part.rs index 3fa8650e..df2e44bc 100644 --- a/src/typing/multipart/part.rs +++ b/src/typing/multipart/part.rs @@ -24,7 +24,7 @@ pub struct Part { } /// The data for a part of a multipart form. -pub enum PartData { +pub enum PartExtractor { Text(Bytes), Bytes(Bytes), File(PathBuf), @@ -33,7 +33,7 @@ pub enum PartData { } #[cfg(feature = "docs")] -impl PyStubType for PartData { +impl PyStubType for PartExtractor { fn type_output() -> TypeInfo { TypeInfo::any() } @@ -43,28 +43,34 @@ impl PyStubType for PartData { #[pymethods] impl Part { /// Creates a new part. + /// + /// # Arguments + /// - `name` - The name of the part. + /// - `value` - The value of the part, either text, bytes, a file path, or a async or sync stream. + /// - `filename` - The filename of the part. + /// - `mime` - The MIME type of the part. #[new] #[pyo3(signature = (name, value, filename = None, mime = None))] pub fn new( py: Python, name: String, - value: PartData, + value: PartExtractor, filename: Option, mime: Option<&str>, ) -> PyResult { py.allow_threads(|| { // Create the inner part let mut inner = match value { - PartData::Text(bytes) | PartData::Bytes(bytes) => { + PartExtractor::Text(bytes) | PartExtractor::Bytes(bytes) => { rquest::multipart::Part::stream(Body::from(bytes)) } - PartData::File(path) => pyo3_async_runtimes::tokio::get_runtime() + PartExtractor::File(path) => pyo3_async_runtimes::tokio::get_runtime() .block_on(rquest::multipart::Part::file(path)) .map_err(Error::from)?, - PartData::SyncStream(stream) => { + PartExtractor::SyncStream(stream) => { rquest::multipart::Part::stream(Body::wrap_stream(stream)) } - PartData::AsyncStream(stream) => { + PartExtractor::AsyncStream(stream) => { rquest::multipart::Part::stream(Body::wrap_stream(stream)) } }; @@ -89,7 +95,7 @@ impl Part { } } -impl FromPyObject<'_> for PartData { +impl FromPyObject<'_> for PartExtractor { fn extract_bound(ob: &Bound<'_, PyAny>) -> PyResult { if let Ok(text) = ob.extract::() { return Ok(Self::Text(Bytes::from_owner(text))); diff --git a/src/typing/param/client.rs b/src/typing/param/client.rs index 958bc200..0c4b33ac 100644 --- a/src/typing/param/client.rs +++ b/src/typing/param/client.rs @@ -1,8 +1,11 @@ -use crate::typing::{ - HeaderMapFromPy, HeadersOrderFromPyList, Impersonate, ImpersonateOS, IpAddr, LookupIpStrategy, - Proxy, SslVerify, TlsVersion, +use crate::{ + extract_option, + typing::{ + HeaderMapExtractor, HeadersOrderExtractor, ImpersonateExtractor, IpAddr, LookupIpStrategy, + SslVerify, TlsVersion, proxy::ProxyListExtractor, + }, }; -use pyo3::{prelude::*, pybacked::PyBackedStr, types::PyList}; +use pyo3::{prelude::*, pybacked::PyBackedStr}; #[cfg(feature = "docs")] use pyo3_stub_gen::{PyStubType, TypeInfo}; @@ -10,25 +13,16 @@ use pyo3_stub_gen::{PyStubType, TypeInfo}; #[derive(Default)] pub struct ClientParams { /// The impersonation settings for the request. - pub impersonate: Option, - - /// The impersonation settings for the operating system. - pub impersonate_os: Option, - - /// Whether to skip impersonate HTTP/2. - pub impersonate_skip_http2: Option, - - /// Whether to skip impersonate headers. - pub impersonate_skip_headers: Option, + pub impersonate: Option, /// The user agent to use for the request. pub user_agent: Option, /// The headers to use for the request. - pub default_headers: Option, + pub default_headers: Option, /// The order of the headers to use for the request. - pub headers_order: Option, + pub headers_order: Option, /// Whether to use referer. pub referer: Option, @@ -104,7 +98,7 @@ pub struct ClientParams { pub no_proxy: Option, /// The proxy to use for the request. - pub proxies: Option>, + pub proxies: Option, /// Bind to a local IP Address. pub local_address: Option, @@ -130,26 +124,17 @@ pub struct ClientParams { #[derive(Default)] pub struct UpdateClientParams { /// The impersonation settings for the request. - pub impersonate: Option, - - /// The impersonation settings for the operating system. - pub impersonate_os: Option, - - /// Whether to skip impersonate HTTP/2. - pub impersonate_skip_http2: Option, - - /// Whether to skip impersonate headers. - pub impersonate_skip_headers: Option, + pub impersonate: Option, /// The headers to use for the request. - pub headers: Option, + pub headers: Option, /// The order of the headers to use for the request. - pub headers_order: Option, + pub headers_order: Option, // ========= Network options ========= /// The proxy to use for the request. - pub proxies: Option>, + pub proxies: Option, /// Bind to a local IP Address. pub local_address: Option, @@ -158,40 +143,10 @@ pub struct UpdateClientParams { pub interface: Option, } -macro_rules! extract_option { - ($ob:expr, $params:expr, $field:ident) => { - if let Ok(value) = $ob.get_item(stringify!($field)) { - $params.$field = value.extract()?; - } - }; -} - -fn extract_proxies(ob: &Bound<'_, PyAny>) -> PyResult>> { - if let Ok(proxies) = ob.get_item("proxies") { - let proxies = proxies.downcast::()?; - let len = proxies.len(); - proxies - .into_iter() - .try_fold(Vec::with_capacity(len), |mut list, proxy| { - let proxy = proxy.downcast::()?; - if let Some(proxy) = proxy.borrow_mut().0.take() { - list.push(proxy); - } - Ok::<_, PyErr>(list) - }) - .map(Some) - } else { - Ok(None) - } -} - impl<'py> FromPyObject<'py> for ClientParams { fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { let mut params = Self::default(); extract_option!(ob, params, impersonate); - extract_option!(ob, params, impersonate_os); - extract_option!(ob, params, impersonate_skip_http2); - extract_option!(ob, params, impersonate_skip_headers); extract_option!(ob, params, user_agent); extract_option!(ob, params, default_headers); @@ -211,7 +166,7 @@ impl<'py> FromPyObject<'py> for ClientParams { extract_option!(ob, params, tcp_keepalive); extract_option!(ob, params, no_proxy); - params.proxies = extract_proxies(ob)?; + extract_option!(ob, params, proxies); extract_option!(ob, params, local_address); extract_option!(ob, params, interface); @@ -237,12 +192,9 @@ impl<'py> FromPyObject<'py> for UpdateClientParams { fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { let mut params = Self::default(); extract_option!(ob, params, impersonate); - extract_option!(ob, params, impersonate_os); - extract_option!(ob, params, impersonate_skip_http2); - extract_option!(ob, params, impersonate_skip_headers); extract_option!(ob, params, headers); extract_option!(ob, params, headers_order); - params.proxies = extract_proxies(ob)?; + extract_option!(ob, params, proxies); extract_option!(ob, params, local_address); extract_option!(ob, params, interface); Ok(params) diff --git a/src/typing/param/mod.rs b/src/typing/param/mod.rs index 7571d6ad..9f55a497 100644 --- a/src/typing/param/mod.rs +++ b/src/typing/param/mod.rs @@ -5,3 +5,12 @@ mod ws; pub use self::client::{ClientParams, UpdateClientParams}; pub use self::request::RequestParams; pub use self::ws::WebSocketParams; + +#[macro_export] +macro_rules! extract_option { + ($ob:expr, $params:expr, $field:ident) => { + if let Ok(value) = $ob.get_item(stringify!($field)) { + $params.$field = value.extract()?; + } + }; +} diff --git a/src/typing/param/request.rs b/src/typing/param/request.rs index 2881c533..f0402f6a 100644 --- a/src/typing/param/request.rs +++ b/src/typing/param/request.rs @@ -1,6 +1,9 @@ -use crate::typing::{ - CookieFromPyDict, FromPyBody, HeaderMapFromPy, IpAddr, Json, Multipart, Proxy, QueryOrForm, - Version, +use crate::{ + extract_option, + typing::{ + BodyExtractor, CookieExtractor, HeaderMapExtractor, IpAddr, Json, ProxyExtractor, + UrlEncodedValuesExtractor, Version, multipart::MultipartExtractor, + }, }; use pyo3::{prelude::*, pybacked::PyBackedStr}; #[cfg(feature = "docs")] @@ -10,7 +13,7 @@ use pyo3_stub_gen::{PyStubType, TypeInfo}; #[derive(Default)] pub struct RequestParams { /// The proxy to use for the request. - pub proxy: Option, + pub proxy: Option, /// Bind to a local IP Address. pub local_address: Option, @@ -28,10 +31,10 @@ pub struct RequestParams { pub version: Option, /// The headers to use for the request. - pub headers: Option, + pub headers: Option, /// The cookies to use for the request. - pub cookies: Option, + pub cookies: Option, /// Whether to allow redirects. pub allow_redirects: Option, @@ -49,36 +52,25 @@ pub struct RequestParams { pub basic_auth: Option<(PyBackedStr, Option)>, /// The query parameters to use for the request. - pub query: Option, + pub query: Option, /// The form parameters to use for the request. - pub form: Option, + pub form: Option, /// The JSON body to use for the request. pub json: Option, /// The body to use for the request. - pub body: Option, + pub body: Option, /// The multipart form to use for the request. - pub multipart: Option, -} - -macro_rules! extract_option { - ($ob:expr, $params:expr, $field:ident) => { - if let Ok(value) = $ob.get_item(stringify!($field)) { - $params.$field = value.extract()?; - } - }; + pub multipart: Option, } impl<'py> FromPyObject<'py> for RequestParams { fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { let mut params = Self::default(); - if let Ok(ob) = ob.get_item("proxy") { - let proxy = ob.downcast::()?; - params.proxy = proxy.borrow_mut().0.take(); - } + extract_option!(ob, params, proxy); extract_option!(ob, params, local_address); extract_option!(ob, params, interface); extract_option!(ob, params, timeout); @@ -96,10 +88,7 @@ impl<'py> FromPyObject<'py> for RequestParams { extract_option!(ob, params, form); extract_option!(ob, params, json); extract_option!(ob, params, body); - if let Ok(value) = ob.get_item("multipart") { - let form = value.downcast::()?; - params.multipart = form.borrow_mut().0.take(); - } + extract_option!(ob, params, multipart); Ok(params) } diff --git a/src/typing/param/ws.rs b/src/typing/param/ws.rs index 8b6fad0e..b3786585 100644 --- a/src/typing/param/ws.rs +++ b/src/typing/param/ws.rs @@ -1,4 +1,7 @@ -use crate::typing::{CookieFromPyDict, HeaderMapFromPy, IpAddr, QueryOrForm}; +use crate::{ + extract_option, + typing::{CookieExtractor, HeaderMapExtractor, IpAddr, UrlEncodedValuesExtractor}, +}; use pyo3::{prelude::*, pybacked::PyBackedStr}; #[cfg(feature = "docs")] use pyo3_stub_gen::{PyStubType, TypeInfo}; @@ -16,10 +19,10 @@ pub struct WebSocketParams { pub interface: Option, /// The headers to use for the request. - pub headers: Option, + pub headers: Option, /// The cookies to use for the request. - pub cookies: Option, + pub cookies: Option, /// The protocols to use for the request. pub protocols: Option>, @@ -37,7 +40,7 @@ pub struct WebSocketParams { pub basic_auth: Option<(PyBackedStr, Option)>, /// The query parameters to use for the request. - pub query: Option, + pub query: Option, /// Read buffer capacity. This buffer is eagerly allocated and used for receiving /// messages. @@ -92,14 +95,6 @@ pub struct WebSocketParams { pub accept_unmasked_frames: Option, } -macro_rules! extract_option { - ($ob:expr, $params:expr, $field:ident) => { - if let Ok(value) = $ob.get_item(stringify!($field)) { - $params.$field = value.extract()?; - } - }; -} - impl<'py> FromPyObject<'py> for WebSocketParams { fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { let mut params = Self::default(); diff --git a/src/typing/proxy.rs b/src/typing/proxy.rs index 8d4d923f..0770923f 100644 --- a/src/typing/proxy.rs +++ b/src/typing/proxy.rs @@ -1,7 +1,7 @@ use crate::error::Error; -use super::HeaderMapFromPy; -use pyo3::prelude::*; +use super::HeaderMapExtractor; +use pyo3::{prelude::*, types::PyList}; #[cfg(feature = "docs")] use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods}; use rquest::header::HeaderValue; @@ -29,7 +29,7 @@ macro_rules! proxy_method { username: Option<&str>, password: Option<&str>, custom_http_auth: Option<&str>, - custom_http_headers: Option, + custom_http_headers: Option, exclusion: Option<&str>, ) -> PyResult { py.allow_threads(|| { @@ -149,7 +149,7 @@ impl Proxy { username: Option<&'a str>, password: Option<&str>, custom_http_auth: Option<&'a str>, - custom_http_headers: Option, + custom_http_headers: Option, exclusion: Option<&'a str>, ) -> PyResult { let mut proxy = proxy_fn(url).map_err(Error::RquestError)?; @@ -176,3 +176,36 @@ impl Proxy { Ok(Proxy(Some(proxy))) } } + +pub struct ProxyExtractor(pub rquest::Proxy); + +impl FromPyObject<'_> for ProxyExtractor { + fn extract_bound(ob: &Bound<'_, PyAny>) -> PyResult { + let proxy = ob.downcast::()?; + proxy + .borrow_mut() + .0 + .take() + .map(ProxyExtractor) + .ok_or_else(|| Error::MemoryError) + .map_err(Into::into) + } +} + +pub struct ProxyListExtractor(pub Vec); +impl FromPyObject<'_> for ProxyListExtractor { + fn extract_bound(ob: &Bound<'_, PyAny>) -> PyResult { + let proxies = ob.downcast::()?; + let len = proxies.len(); + proxies + .into_iter() + .try_fold(Vec::with_capacity(len), |mut list, proxy| { + let proxy = proxy.downcast::()?; + if let Some(proxy) = proxy.borrow_mut().0.take() { + list.push(proxy); + } + Ok::<_, PyErr>(list) + }) + .map(Self) + } +} From b9e4f371154931799f78cbe0598817629581126c Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Thu, 27 Mar 2025 10:55:39 +0800 Subject: [PATCH 2/2] update tests --- tests/client_test.py | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/tests/client_test.py b/tests/client_test.py index 98852381..0db304b9 100644 --- a/tests/client_test.py +++ b/tests/client_test.py @@ -1,6 +1,6 @@ import pytest import rnet -from rnet import Cookie, Impersonate, ImpersonateOS, HeaderMap +from rnet import Cookie, Impersonate, ImpersonateOS, ImpersonateOption, HeaderMap @pytest.mark.asyncio @@ -19,9 +19,11 @@ def __init__(self, **kwargs): assert client.cookie_jar is None assert client.test_var == "test" client.update( - impersonate=Impersonate.Firefox135, - impersonate_os=ImpersonateOS.Windows, - Impersonate_skip_headers=False, + impersonate=ImpersonateOption( + impersonate=Impersonate.Firefox135, + impersonate_os=ImpersonateOS.Windows, + skip_headers=False, + ) ) assert ( client.user_agent @@ -66,9 +68,11 @@ async def test_update_impersonate(): == "Mozilla/5.0 (Macintosh; Intel Mac OS X 10.15; rv:133.0) Gecko/20100101 Firefox/133.0" ) client.update( - impersonate=Impersonate.Firefox135, - impersonate_os=ImpersonateOS.Windows, - Impersonate_skip_headers=False, + impersonate=ImpersonateOption( + impersonate=Impersonate.Firefox135, + impersonate_os=ImpersonateOS.Windows, + skip_headers=False, + ) ) assert ( client.user_agent