diff --git a/rnet.pyi b/rnet.pyi index 5b90090d..4b741666 100644 --- a/rnet.pyi +++ b/rnet.pyi @@ -1399,9 +1399,19 @@ class HeaderMap: def __contains__(self, key: str) -> builtins.bool: ... def __len__(self) -> builtins.int: ... def __iter__(self) -> HeaderMapKeysIter: ... - def items(self) -> HeaderMapItemsIter: ... def __str__(self) -> builtins.str: ... def __repr__(self) -> builtins.str: ... + def get_all(self, key: str) -> HeaderMapValuesIter: + r""" + Returns multiple value sequences of key mapping + """ + ... + + def items(self) -> HeaderMapItemsIter: + r""" + Returns key-value pairs in the order they were added. + """ + ... class HeaderMapItemsIter: r""" @@ -1421,6 +1431,14 @@ class HeaderMapKeysIter: def __iter__(self) -> HeaderMapKeysIter: ... def __next__(self) -> typing.Optional[typing.Any]: ... +class HeaderMapValuesIter: + r""" + An iterator over the values in a HeaderMap. + """ + + def __iter__(self) -> HeaderMapValuesIter: ... + def __next__(self) -> typing.Optional[typing.Any]: ... + class ImpersonateOption: r""" A struct to represent the `ImpersonateOption` class. diff --git a/src/lib.rs b/src/lib.rs index fd90d29b..262a314a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -14,9 +14,9 @@ use pyo3_async_runtimes::tokio::future_into_py; use pyo3_stub_gen::{define_stub_info_gatherer, derive::*}; use typing::param::{RequestParams, WebSocketParams}; use typing::{ - Cookie, HeaderMap, HeaderMapItemsIter, HeaderMapKeysIter, Impersonate, ImpersonateOS, - ImpersonateOption, LookupIpStrategy, Method, Multipart, Part, Proxy, SameSite, SocketAddr, - StatusCode, TlsVersion, Version, + Cookie, HeaderMap, HeaderMapItemsIter, HeaderMapKeysIter, HeaderMapValuesIter, Impersonate, + ImpersonateOS, ImpersonateOption, LookupIpStrategy, Method, Multipart, Part, Proxy, SameSite, + SocketAddr, StatusCode, TlsVersion, Version, }; #[cfg(not(target_env = "msvc"))] @@ -323,6 +323,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/headers.rs b/src/typing/headers.rs index fc017836..89687f78 100644 --- a/src/typing/headers.rs +++ b/src/typing/headers.rs @@ -46,6 +46,30 @@ impl HeaderMap { Self(headers) } + /// Returns multiple value sequences of key mapping + fn get_all(&self, key: PyBackedStr) -> HeaderMapValuesIter { + HeaderMapValuesIter { + inner: self + .0 + .get_all(key.as_ref() as &str) + .iter() + .cloned() + .collect(), + } + } + + /// Returns key-value pairs in the order they were added. + #[inline] + fn items(&self) -> HeaderMapItemsIter { + HeaderMapItemsIter { + inner: self.0.iter().map(|(k, v)| (k.clone(), v.clone())).collect(), + } + } +} + +#[cfg_attr(feature = "docs", gen_stub_pymethods)] +#[pymethods] +impl HeaderMap { #[inline] fn __getitem__<'py>(&self, py: Python<'py>, key: PyBackedStr) -> Option> { let value = self.0.get(key.as_ref() as &str)?; @@ -89,13 +113,6 @@ impl HeaderMap { } } - #[inline] - fn items(&self) -> HeaderMapItemsIter { - HeaderMapItemsIter { - inner: self.0.iter().map(|(k, v)| (k.clone(), v.clone())).collect(), - } - } - #[inline] fn __str__(&self) -> String { format!("{:?}", self.0) @@ -130,6 +147,28 @@ impl HeaderMapKeysIter { } } +/// An iterator over the values in a HeaderMap. +#[cfg_attr(feature = "docs", gen_stub_pyclass)] +#[pyclass] +pub struct HeaderMapValuesIter { + inner: Vec, +} +#[cfg_attr(feature = "docs", gen_stub_pymethods)] +#[pymethods] +impl HeaderMapValuesIter { + #[inline] + fn __iter__(slf: PyRefMut<'_, Self>) -> PyRefMut<'_, Self> { + slf + } + + #[inline] + fn __next__(mut slf: PyRefMut) -> Option> { + slf.inner + .pop() + .and_then(|v| HeaderValueBuffer::new(v).into_bytes_ref(slf.py()).ok()) + } +} + /// An iterator over the items in a HeaderMap. #[cfg_attr(feature = "docs", gen_stub_pyclass)] #[pyclass] diff --git a/src/typing/mod.rs b/src/typing/mod.rs index 12532bf4..c1d5d06a 100644 --- a/src/typing/mod.rs +++ b/src/typing/mod.rs @@ -15,7 +15,8 @@ pub use self::{ cookie::{Cookie, CookieExtractor}, enums::{Impersonate, ImpersonateOS, LookupIpStrategy, Method, SameSite, TlsVersion, Version}, headers::{ - HeaderMap, HeaderMapExtractor, HeaderMapItemsIter, HeaderMapKeysIter, HeadersOrderExtractor, + HeaderMap, HeaderMapExtractor, HeaderMapItemsIter, HeaderMapKeysIter, HeaderMapValuesIter, + HeadersOrderExtractor, }, ipaddr::{IpAddr, SocketAddr}, json::Json,