From 7887086de5c6320dbddf3488e583302ed40001b1 Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Tue, 25 Mar 2025 08:02:32 +0800 Subject: [PATCH 1/2] feat(headers): adapt HeaderMap to support repeated headers --- rnet.pyi | 10 +++++-- src/typing/headers.rs | 55 +++++++++++++++++++++++++++++++++---- src/typing/mod.rs | 3 +- src/typing/param/client.rs | 8 +++--- src/typing/param/request.rs | 4 +-- src/typing/param/ws.rs | 4 +-- src/typing/proxy.rs | 10 +++---- tests/client_test.py | 5 +++- tests/request_test.py | 9 ++++-- 9 files changed, 81 insertions(+), 27 deletions(-) diff --git a/rnet.pyi b/rnet.pyi index e47fa546..ca1053dc 100644 --- a/rnet.pyi +++ b/rnet.pyi @@ -1405,12 +1405,16 @@ class HeaderMap: r""" A HTTP header map. """ + def __new__(cls,init:typing.Optional[dict]): ... def __getitem__(self, key:str) -> typing.Optional[typing.Any]: ... def __setitem__(self, key:str, value:str) -> None: ... + def __append__(self, key:str, value:str) -> None: + ... + def __delitem__(self, key:str) -> None: ... @@ -1606,7 +1610,7 @@ class Proxy: Supports HTTP, HTTPS, SOCKS4, SOCKS4a, SOCKS5, and SOCKS5h protocols. """ @staticmethod - def http(url:builtins.str, username:typing.Optional[builtins.str]=None, password:typing.Optional[builtins.str]=None, custom_http_auth:typing.Optional[builtins.str]=None, custom_httt_headers:typing.Optional[typing.Dict[str, str]]=None, exclusion:typing.Optional[builtins.str]=None) -> Proxy: + def http(url:builtins.str, username:typing.Optional[builtins.str]=None, password:typing.Optional[builtins.str]=None, custom_http_auth:typing.Optional[builtins.str]=None, custom_httt_headers:typing.Optional[typing.Union[typing.Dict[str, str], HeaderMap]]=None, exclusion:typing.Optional[builtins.str]=None) -> Proxy: r""" Creates a new HTTP proxy. @@ -1635,7 +1639,7 @@ class Proxy: ... @staticmethod - def https(url:builtins.str, username:typing.Optional[builtins.str]=None, password:typing.Optional[builtins.str]=None, custom_http_auth:typing.Optional[builtins.str]=None, custom_httt_headers:typing.Optional[typing.Dict[str, str]]=None, exclusion:typing.Optional[builtins.str]=None) -> Proxy: + def https(url:builtins.str, username:typing.Optional[builtins.str]=None, password:typing.Optional[builtins.str]=None, custom_http_auth:typing.Optional[builtins.str]=None, custom_httt_headers:typing.Optional[typing.Union[typing.Dict[str, str], HeaderMap]]=None, exclusion:typing.Optional[builtins.str]=None) -> Proxy: r""" Creates a new HTTPS proxy. @@ -1664,7 +1668,7 @@ class Proxy: ... @staticmethod - def all(url:builtins.str, username:typing.Optional[builtins.str]=None, password:typing.Optional[builtins.str]=None, custom_http_auth:typing.Optional[builtins.str]=None, custom_httt_headers:typing.Optional[typing.Dict[str, str]]=None, exclusion:typing.Optional[builtins.str]=None) -> Proxy: + def all(url:builtins.str, username:typing.Optional[builtins.str]=None, password:typing.Optional[builtins.str]=None, custom_http_auth:typing.Optional[builtins.str]=None, custom_httt_headers:typing.Optional[typing.Union[typing.Dict[str, str], HeaderMap]]=None, exclusion:typing.Optional[builtins.str]=None) -> Proxy: r""" Creates a new proxy for all protocols. diff --git a/src/typing/headers.rs b/src/typing/headers.rs index ce2dbd70..8fe0251b 100644 --- a/src/typing/headers.rs +++ b/src/typing/headers.rs @@ -22,6 +22,30 @@ pub struct HeaderMap(pub header::HeaderMap); #[cfg_attr(feature = "docs", gen_stub_pymethods)] #[pymethods] impl HeaderMap { + #[new] + #[inline] + fn new(init: Option<&Bound<'_, PyDict>>) -> Self { + let mut headers = header::HeaderMap::new(); + + // This section of memory might be retained by the Rust object, + // and we want to prevent Python's garbage collector from managing it. + if let Some(dict) = init { + for (name, value) in dict.iter() { + if let (Ok(Ok(name)), Ok(Ok(value))) = ( + name.extract::() + .map(|n| HeaderName::from_bytes(n.as_bytes())), + value + .extract::() + .map(|v| HeaderValue::from_bytes(v.as_bytes())), + ) { + headers.insert(name, value); + } + } + } + + Self(headers) + } + #[inline] fn __getitem__<'py>(&self, py: Python<'py>, key: PyBackedStr) -> Option> { let value = self.0.get(key.as_ref() as &str)?; @@ -41,6 +65,18 @@ impl HeaderMap { }) } + #[inline] + fn __append__(&mut self, py: Python, key: PyBackedStr, value: PyBackedStr) { + py.allow_threads(|| { + if let (Ok(name), Ok(value)) = ( + HeaderName::from_bytes(key.as_bytes()), + HeaderValue::from_bytes(value.as_bytes()), + ) { + self.0.append(name, value); + } + }) + } + #[inline] fn __delitem__(&mut self, py: Python, key: PyBackedStr) { py.allow_threads(|| { @@ -135,15 +171,18 @@ impl HeaderMapItemsIter { } /// A HTTP header map. -pub struct HeaderMapFromPyDict(pub header::HeaderMap); +pub struct HeaderMapFromPy(pub header::HeaderMap); /// A list of header names in order. pub struct HeadersOrderFromPyList(pub Vec); -impl FromPyObject<'_> for HeaderMapFromPyDict { +impl FromPyObject<'_> for HeaderMapFromPy { fn extract_bound(ob: &Bound<'_, PyAny>) -> PyResult { - let dict = ob.downcast::()?; + if let Ok(headers) = ob.downcast::() { + return Ok(Self(headers.borrow().0.clone())); + } + let dict = ob.downcast::()?; dict.iter() .try_fold( header::HeaderMap::with_capacity(dict.len()), @@ -160,7 +199,8 @@ impl FromPyObject<'_> for HeaderMapFromPyDict { } } -impl<'py> IntoPyObject<'py> for HeaderMapFromPyDict { +#[cfg(feature = "docs")] +impl<'py> IntoPyObject<'py> for HeaderMapFromPy { type Target = HeaderMap; type Output = Bound<'py, Self::Target>; @@ -173,9 +213,12 @@ impl<'py> IntoPyObject<'py> for HeaderMapFromPyDict { } #[cfg(feature = "docs")] -impl PyStubType for HeaderMapFromPyDict { +impl PyStubType for HeaderMapFromPy { fn type_output() -> TypeInfo { - TypeInfo::with_module("typing.Dict[str, str]", "typing".into()) + TypeInfo::with_module( + "typing.Union[typing.Dict[str, str], HeaderMap]", + "typing".into(), + ) } } diff --git a/src/typing/mod.rs b/src/typing/mod.rs index 45933c83..c9afa865 100644 --- a/src/typing/mod.rs +++ b/src/typing/mod.rs @@ -15,8 +15,7 @@ pub use self::{ cookie::{Cookie, CookieFromPyDict}, enums::{Impersonate, ImpersonateOS, LookupIpStrategy, Method, SameSite, TlsVersion, Version}, headers::{ - HeaderMap, HeaderMapFromPyDict, HeaderMapItemsIter, HeaderMapKeysIter, - HeadersOrderFromPyList, + HeaderMap, HeaderMapFromPy, HeaderMapItemsIter, HeaderMapKeysIter, HeadersOrderFromPyList, }, ipaddr::{IpAddr, SocketAddr}, json::Json, diff --git a/src/typing/param/client.rs b/src/typing/param/client.rs index 3a06bd1c..958bc200 100644 --- a/src/typing/param/client.rs +++ b/src/typing/param/client.rs @@ -1,6 +1,6 @@ use crate::typing::{ - HeaderMapFromPyDict, HeadersOrderFromPyList, Impersonate, ImpersonateOS, IpAddr, - LookupIpStrategy, Proxy, SslVerify, TlsVersion, + HeaderMapFromPy, HeadersOrderFromPyList, Impersonate, ImpersonateOS, IpAddr, LookupIpStrategy, + Proxy, SslVerify, TlsVersion, }; use pyo3::{prelude::*, pybacked::PyBackedStr, types::PyList}; #[cfg(feature = "docs")] @@ -25,7 +25,7 @@ pub struct ClientParams { 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, @@ -142,7 +142,7 @@ pub struct UpdateClientParams { pub impersonate_skip_headers: 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, diff --git a/src/typing/param/request.rs b/src/typing/param/request.rs index e003327e..2881c533 100644 --- a/src/typing/param/request.rs +++ b/src/typing/param/request.rs @@ -1,5 +1,5 @@ use crate::typing::{ - CookieFromPyDict, FromPyBody, HeaderMapFromPyDict, IpAddr, Json, Multipart, Proxy, QueryOrForm, + CookieFromPyDict, FromPyBody, HeaderMapFromPy, IpAddr, Json, Multipart, Proxy, QueryOrForm, Version, }; use pyo3::{prelude::*, pybacked::PyBackedStr}; @@ -28,7 +28,7 @@ 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, diff --git a/src/typing/param/ws.rs b/src/typing/param/ws.rs index a6d82f1b..8b6fad0e 100644 --- a/src/typing/param/ws.rs +++ b/src/typing/param/ws.rs @@ -1,4 +1,4 @@ -use crate::typing::{CookieFromPyDict, HeaderMapFromPyDict, IpAddr, QueryOrForm}; +use crate::typing::{CookieFromPyDict, HeaderMapFromPy, IpAddr, QueryOrForm}; use pyo3::{prelude::*, pybacked::PyBackedStr}; #[cfg(feature = "docs")] use pyo3_stub_gen::{PyStubType, TypeInfo}; @@ -16,7 +16,7 @@ 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, diff --git a/src/typing/proxy.rs b/src/typing/proxy.rs index da1d42e3..82145e7b 100644 --- a/src/typing/proxy.rs +++ b/src/typing/proxy.rs @@ -1,6 +1,6 @@ use crate::error::Error; -use super::HeaderMapFromPyDict; +use super::HeaderMapFromPy; use pyo3::prelude::*; #[cfg(feature = "docs")] use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods}; @@ -53,7 +53,7 @@ impl Proxy { username: Option<&str>, password: Option<&str>, custom_http_auth: Option<&str>, - custom_httt_headers: Option, + custom_httt_headers: Option, exclusion: Option<&str>, ) -> PyResult { Self::create_proxy( @@ -105,7 +105,7 @@ impl Proxy { username: Option<&str>, password: Option<&str>, custom_http_auth: Option<&str>, - custom_httt_headers: Option, + custom_httt_headers: Option, exclusion: Option<&str>, ) -> PyResult { Self::create_proxy( @@ -157,7 +157,7 @@ impl Proxy { username: Option<&str>, password: Option<&str>, custom_http_auth: Option<&str>, - custom_httt_headers: Option, + custom_httt_headers: Option, exclusion: Option<&str>, ) -> PyResult { Self::create_proxy( @@ -179,7 +179,7 @@ impl Proxy { username: Option<&'a str>, password: Option<&str>, custom_http_auth: Option<&'a str>, - custom_httt_headers: Option, + custom_httt_headers: Option, exclusion: Option<&'a str>, ) -> PyResult { let mut proxy = proxy_fn(url).map_err(Error::RquestError)?; diff --git a/tests/client_test.py b/tests/client_test.py index 97c482ab..98852381 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 +from rnet import Cookie, Impersonate, ImpersonateOS, HeaderMap @pytest.mark.asyncio @@ -38,6 +38,9 @@ async def test_update_headers(): client.update(headers=headers) assert client.headers["user-agent"] == b"rnet" + client.update(headers=HeaderMap(headers)) + assert client.headers["user-agent"] == b"rnet" + @pytest.mark.asyncio @pytest.mark.flaky(reruns=3, reruns_delay=2) diff --git a/tests/request_test.py b/tests/request_test.py index 0d7a74fd..3fd34dfc 100644 --- a/tests/request_test.py +++ b/tests/request_test.py @@ -1,6 +1,6 @@ import pytest import rnet -from rnet import Version +from rnet import Version, HeaderMap client = rnet.Client(tls_info=True) @@ -20,7 +20,12 @@ async def test_send_with_version(): @pytest.mark.flaky(reruns=3, reruns_delay=2) async def test_send_headers(): url = "https://httpbin.org/headers" - response = await client.get(url, headers={"foo": "bar"}) + headers = {"foo": "bar"} + response = await client.get(url, headers=headers) + json = await response.json() + assert json["headers"]["Foo"] == "bar" + + response = await client.get(url, headers=HeaderMap(headers)) json = await response.json() assert json["headers"]["Foo"] == "bar" From 4a065d2bd2bd88bdcdf2e012d5d89aea3cb1c6c3 Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Tue, 25 Mar 2025 08:07:07 +0800 Subject: [PATCH 2/2] feat(headers): adapt HeaderMap to support repeated headers --- rnet.pyi | 3 --- src/typing/headers.rs | 12 ------------ 2 files changed, 15 deletions(-) diff --git a/rnet.pyi b/rnet.pyi index ca1053dc..d976dd62 100644 --- a/rnet.pyi +++ b/rnet.pyi @@ -1412,9 +1412,6 @@ class HeaderMap: def __setitem__(self, key:str, value:str) -> None: ... - def __append__(self, key:str, value:str) -> None: - ... - def __delitem__(self, key:str) -> None: ... diff --git a/src/typing/headers.rs b/src/typing/headers.rs index 8fe0251b..e1babb75 100644 --- a/src/typing/headers.rs +++ b/src/typing/headers.rs @@ -55,18 +55,6 @@ impl HeaderMap { #[inline] fn __setitem__(&mut self, py: Python, key: PyBackedStr, value: PyBackedStr) { - py.allow_threads(|| { - if let (Ok(name), Ok(value)) = ( - HeaderName::from_bytes(key.as_bytes()), - HeaderValue::from_bytes(value.as_bytes()), - ) { - self.0.insert(name, value); - } - }) - } - - #[inline] - fn __append__(&mut self, py: Python, key: PyBackedStr, value: PyBackedStr) { py.allow_threads(|| { if let (Ok(name), Ok(value)) = ( HeaderName::from_bytes(key.as_bytes()),