feat(headers): adapt HeaderMap to support repeated headers

rnet.pyi

@@ -1405,6 +1405,7 @@ class HeaderMap:
     r"""
     A HTTP header map.
     """
+    def __new__(cls,init:typing.Optional[dict]): ...
     def __getitem__(self, key:str) -> typing.Optional[typing.Any]:
         ...
 
@@ -1606,7 +1607,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 +1636,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 +1665,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.
 

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::<PyBackedStr>()
+                        .map(|n| HeaderName::from_bytes(n.as_bytes())),
+                    value
+                        .extract::<PyBackedStr>()
+                        .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<Bound<'py, PyAny>> {
         let value = self.0.get(key.as_ref() as &str)?;
@@ -36,7 +60,7 @@ impl HeaderMap {
                 HeaderName::from_bytes(key.as_bytes()),
                 HeaderValue::from_bytes(value.as_bytes()),
             ) {
-                self.0.insert(name, value);
+                self.0.append(name, value);
             }
         })
     }
@@ -135,15 +159,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<HeaderName>);
 
-impl FromPyObject<'_> for HeaderMapFromPyDict {
+impl FromPyObject<'_> for HeaderMapFromPy {
     fn extract_bound(ob: &Bound<'_, PyAny>) -> PyResult<Self> {
-        let dict = ob.downcast::<PyDict>()?;
+        if let Ok(headers) = ob.downcast::<HeaderMap>() {
+            return Ok(Self(headers.borrow().0.clone()));
+        }
 
+        let dict = ob.downcast::<PyDict>()?;
         dict.iter()
             .try_fold(
                 header::HeaderMap::with_capacity(dict.len()),
@@ -160,7 +187,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 +201,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(),
+        )
     }
 }
 

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,

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<PyBackedStr>,
 
     /// The headers to use for the request.
-    pub default_headers: Option<HeaderMapFromPyDict>,
+    pub default_headers: Option<HeaderMapFromPy>,
 
     /// The order of the headers to use for the request.
     pub headers_order: Option<HeadersOrderFromPyList>,
@@ -142,7 +142,7 @@ pub struct UpdateClientParams {
     pub impersonate_skip_headers: Option<bool>,
 
     /// The headers to use for the request.
-    pub headers: Option<HeaderMapFromPyDict>,
+    pub headers: Option<HeaderMapFromPy>,
 
     /// The order of the headers to use for the request.
     pub headers_order: Option<HeadersOrderFromPyList>,

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<Version>,
 
     /// The headers to use for the request.
-    pub headers: Option<HeaderMapFromPyDict>,
+    pub headers: Option<HeaderMapFromPy>,
 
     /// The cookies to use for the request.
     pub cookies: Option<CookieFromPyDict>,

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<String>,
 
     /// The headers to use for the request.
-    pub headers: Option<HeaderMapFromPyDict>,
+    pub headers: Option<HeaderMapFromPy>,
 
     /// The cookies to use for the request.
     pub cookies: Option<CookieFromPyDict>,

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<HeaderMapFromPyDict>,
+        custom_httt_headers: Option<HeaderMapFromPy>,
         exclusion: Option<&str>,
     ) -> PyResult<Self> {
         Self::create_proxy(
@@ -105,7 +105,7 @@ impl Proxy {
         username: Option<&str>,
         password: Option<&str>,
         custom_http_auth: Option<&str>,
-        custom_httt_headers: Option<HeaderMapFromPyDict>,
+        custom_httt_headers: Option<HeaderMapFromPy>,
         exclusion: Option<&str>,
     ) -> PyResult<Self> {
         Self::create_proxy(
@@ -157,7 +157,7 @@ impl Proxy {
         username: Option<&str>,
         password: Option<&str>,
         custom_http_auth: Option<&str>,
-        custom_httt_headers: Option<HeaderMapFromPyDict>,
+        custom_httt_headers: Option<HeaderMapFromPy>,
         exclusion: Option<&str>,
     ) -> PyResult<Self> {
         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<HeaderMapFromPyDict>,
+        custom_httt_headers: Option<HeaderMapFromPy>,
         exclusion: Option<&'a str>,
     ) -> PyResult<Self> {
         let mut proxy = proxy_fn(url).map_err(Error::RquestError)?;

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)

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"