From 00bc7165436ae929db09e9cdc51ec323508a1d89 Mon Sep 17 00:00:00 2001 From: 0x676e67 Date: Mon, 11 Aug 2025 11:29:44 +0800 Subject: [PATCH] refactor(typing): simplify Python class `__str__` magic method implementation --- python/rnet/__init__.pyi | 3 - python/rnet/cookie.pyi | 1 - python/rnet/header.pyi | 4 -- src/client/async_impl/response/http.rs | 5 +- src/client/async_impl/response/ws/message.rs | 14 ++-- src/client/async_impl/response/ws/mod.rs | 3 +- src/client/blocking/response/http.rs | 3 +- src/client/blocking/response/ws.rs | 3 +- src/client/net.rs | 16 ++--- src/client/param/client.rs | 2 +- src/emulation.rs | 2 +- src/{ => http}/cookie.rs | 21 +++--- src/http/header.rs | 71 +++++++++++--------- src/http/mod.rs | 5 +- src/http/status.rs | 21 +++--- src/lib.rs | 8 +-- tests/header_test.py | 12 ++++ 17 files changed, 98 insertions(+), 96 deletions(-) rename src/{ => http}/cookie.rs (97%) diff --git a/python/rnet/__init__.pyi b/python/rnet/__init__.pyi index 473fb6bf..d6870e14 100644 --- a/python/rnet/__init__.pyi +++ b/python/rnet/__init__.pyi @@ -575,7 +575,6 @@ class SocketAddr: """ def __str__(self) -> str: ... - def __repr__(self) -> str: ... def ip(self) -> Union[ipaddress.IPv4Address, ipaddress.IPv6Address]: r""" Returns the IP address of the socket address. @@ -592,7 +591,6 @@ class StatusCode: """ def __str__(self) -> str: ... - def __repr__(self) -> str: ... def as_int(self) -> int: r""" Return the status code as an integer. @@ -1002,7 +1000,6 @@ class Message: Returns the close code and reason of the message if it is a close message. """ def __str__(self) -> str: ... - def __repr__(self) -> str: ... @staticmethod def text_from_json(json: Dict[str, Any]) -> Message: r""" diff --git a/python/rnet/cookie.pyi b/python/rnet/cookie.pyi index 92d021e9..e0d10186 100644 --- a/python/rnet/cookie.pyi +++ b/python/rnet/cookie.pyi @@ -81,7 +81,6 @@ class Cookie: """ def __str__(self) -> str: ... - def __repr__(self) -> str: ... class Jar: r""" diff --git a/python/rnet/header.pyi b/python/rnet/header.pyi index 421af43a..b30b60ff 100644 --- a/python/rnet/header.pyi +++ b/python/rnet/header.pyi @@ -54,10 +54,6 @@ class HeaderMap: """Return a string representation of all headers.""" ... - def __repr__(self) -> str: - """Return a detailed string representation.""" - ... - def __new__( cls, init: Optional[dict] = None, capacity: Optional[int] = None ) -> HeaderMap: diff --git a/src/client/async_impl/response/http.rs b/src/client/async_impl/response/http.rs index 40fade65..77b6e25e 100644 --- a/src/client/async_impl/response/http.rs +++ b/src/client/async_impl/response/http.rs @@ -11,9 +11,8 @@ use wreq::{Url, header, tls::TlsInfo}; use crate::{ buffer::{Buffer, BytesBuffer, PyBufferProtocol}, client::{SocketAddr, json::Json}, - cookie::Cookie, error::Error, - http::{StatusCode, Version, header::HeaderMap}, + http::{Version, cookie::Cookie, header::HeaderMap, status::StatusCode}, }; /// A response from a request. @@ -224,8 +223,8 @@ type InnerStreamer = Pin> + Sen /// Used to stream response content. /// Implemented in the `stream` method of the `Response` class. /// Can be used in an asynchronous for loop in Python. -#[pyclass(subclass)] #[derive(Clone)] +#[pyclass(subclass)] pub struct Streamer(Arc>>); impl Deref for Streamer { diff --git a/src/client/async_impl/response/ws/message.rs b/src/client/async_impl/response/ws/message.rs index 305b2bc9..5b63f305 100644 --- a/src/client/async_impl/response/ws/message.rs +++ b/src/client/async_impl/response/ws/message.rs @@ -1,3 +1,5 @@ +use std::fmt; + use bytes::Bytes; use pyo3::{ prelude::*, @@ -12,8 +14,8 @@ use crate::{ }; /// A WebSocket message. -#[pyclass(subclass)] #[derive(Clone)] +#[pyclass(subclass, str)] pub struct Message(pub message::Message); #[pymethods] @@ -164,12 +166,10 @@ impl Message { }); Message(msg) } +} - fn __str__(&self) -> String { - format!("{:?}", self.0) - } - - fn __repr__(&self) -> String { - self.__str__() +impl fmt::Display for Message { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self.0) } } diff --git a/src/client/async_impl/response/ws/mod.rs b/src/client/async_impl/response/ws/mod.rs index 74ba9229..1e43e040 100644 --- a/src/client/async_impl/response/ws/mod.rs +++ b/src/client/async_impl/response/ws/mod.rs @@ -18,9 +18,8 @@ use wreq::{ use crate::{ client::SocketAddr, - cookie::Cookie, error::Error, - http::{StatusCode, Version, header::HeaderMap}, + http::{Version, cookie::Cookie, header::HeaderMap, status::StatusCode}, }; type Sender = Arc>>>; diff --git a/src/client/blocking/response/http.rs b/src/client/blocking/response/http.rs index 1706eb64..e87f6959 100644 --- a/src/client/blocking/response/http.rs +++ b/src/client/blocking/response/http.rs @@ -9,9 +9,8 @@ use crate::{ async_impl::response::{Response, Streamer}, json::Json, }, - cookie::Cookie, error::Error, - http::{StatusCode, Version, header::HeaderMap}, + http::{Version, cookie::Cookie, header::HeaderMap, status::StatusCode}, }; /// A blocking response from a request. diff --git a/src/client/blocking/response/ws.rs b/src/client/blocking/response/ws.rs index 243ba644..44576dba 100644 --- a/src/client/blocking/response/ws.rs +++ b/src/client/blocking/response/ws.rs @@ -5,9 +5,8 @@ use crate::{ SocketAddr, async_impl::response::{Message, WebSocket}, }, - cookie::Cookie, error::Error, - http::{StatusCode, Version, header::HeaderMap}, + http::{Version, cookie::Cookie, header::HeaderMap, status::StatusCode}, }; /// A blocking WebSocket response. diff --git a/src/client/net.rs b/src/client/net.rs index fbf5c81f..8fdddc7f 100644 --- a/src/client/net.rs +++ b/src/client/net.rs @@ -1,8 +1,10 @@ +use std::fmt; + use pyo3::{IntoPyObjectExt, prelude::*}; /// A IP socket address. -#[pyclass(eq)] -#[derive(Clone, Copy, PartialEq, Eq, Hash, Debug)] +#[derive(Clone, Copy, PartialEq, Eq)] +#[pyclass(eq, str)] pub struct SocketAddr(pub std::net::SocketAddr); #[pymethods] @@ -16,12 +18,10 @@ impl SocketAddr { fn port(&self) -> u16 { self.0.port() } +} - fn __str__(&self) -> String { - self.0.to_string() - } - - fn __repr__(&self) -> String { - self.__str__() +impl fmt::Display for SocketAddr { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.0) } } diff --git a/src/client/param/client.rs b/src/client/param/client.rs index 46fd4707..2694a409 100644 --- a/src/client/param/client.rs +++ b/src/client/param/client.rs @@ -5,8 +5,8 @@ use wreq::{Proxy, header::HeaderMap}; use wreq_util::EmulationOption; use crate::{ - cookie::Jar, extractor::Extractor, + http::cookie::Jar, tls::{SslVerify, TlsVersion}, }; diff --git a/src/emulation.rs b/src/emulation.rs index e4e79ad6..3b58f274 100644 --- a/src/emulation.rs +++ b/src/emulation.rs @@ -95,8 +95,8 @@ define_enum!( ); /// A struct to represent the `EmulationOption` class. -#[pyclass(subclass)] #[derive(Clone)] +#[pyclass(subclass)] pub struct EmulationOption(pub wreq_util::EmulationOption); #[pymethods] diff --git a/src/cookie.rs b/src/http/cookie.rs similarity index 97% rename from src/cookie.rs rename to src/http/cookie.rs index 967f118a..e75e4484 100644 --- a/src/cookie.rs +++ b/src/http/cookie.rs @@ -1,4 +1,5 @@ use std::{ + fmt, sync::{Arc, RwLock}, time::SystemTime, }; @@ -29,8 +30,9 @@ define_enum!( ); /// A single HTTP cookie. -#[pyclass(subclass)] + #[derive(Clone)] +#[pyclass(subclass, str)] pub struct Cookie(pub RawCookie<'static>); /// A good default `CookieStore` implementation. @@ -38,8 +40,8 @@ pub struct Cookie(pub RawCookie<'static>); /// This is the implementation used when simply calling `cookie_store(true)`. /// This type is exposed to allow creating one and filling it with some /// existing cookies more easily, before creating a `Client`. -#[pyclass(subclass)] #[derive(Clone, Default)] +#[pyclass(subclass)] pub struct Jar(Arc>); // ===== impl Cookie ===== @@ -177,17 +179,10 @@ impl Cookie { None | Some(Expiration::Session) => None, } } - - fn __str__(&self) -> String { - self.0.to_string() - } - - fn __repr__(&self) -> String { - self.__str__() - } } impl Cookie { + /// Parse cookies from a `HeaderMap`. pub fn extract_headers_cookies(headers: &HeaderMap) -> Vec { headers .get_all(header::SET_COOKIE) @@ -206,6 +201,12 @@ impl Cookie { } } +impl fmt::Display for Cookie { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.0) + } +} + // ===== impl Jar ===== impl CookieStore for Jar { diff --git a/src/http/header.rs b/src/http/header.rs index 4b3f55c1..97683b1f 100644 --- a/src/http/header.rs +++ b/src/http/header.rs @@ -1,3 +1,5 @@ +use std::fmt; + use pyo3::{ prelude::*, pybacked::{PyBackedBytes, PyBackedStr}, @@ -8,32 +10,43 @@ use wreq::header::{self, HeaderName, HeaderValue}; use crate::buffer::{HeaderNameBuffer, HeaderValueBuffer, PyBufferProtocol}; /// A HTTP header map. -#[pyclass(subclass)] +#[pyclass(subclass, str)] #[derive(Clone)] pub struct HeaderMap(pub header::HeaderMap); #[pymethods] impl HeaderMap { + /// Creates a new `HeaderMap` from an optional dictionary. #[new] - #[pyo3(signature = (init=None, capacity=None))] - fn new(init: Option<&Bound<'_, PyDict>>, capacity: Option) -> Self { + #[pyo3(signature = (dict=None, capacity=None))] + fn new(dict: Option<&Bound<'_, PyDict>>, capacity: Option) -> Self { let mut headers = capacity .map(header::HeaderMap::with_capacity) .unwrap_or_default(); // 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 { + if let Some(dict) = dict { 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(HeaderValue::from_maybe_shared), - ) { - headers.insert(name, value); - } + let name = match name + .extract::() + .ok() + .and_then(|n| HeaderName::from_bytes(n.as_bytes()).ok()) + { + Some(n) => n, + None => continue, + }; + + let value = match value + .extract::() + .ok() + .and_then(|v| HeaderValue::from_maybe_shared(v).ok()) + { + Some(v) => v, + None => continue, + }; + + headers.insert(name, value); } } @@ -52,16 +65,13 @@ impl HeaderMap { key: PyBackedStr, default: Option, ) -> Option> { - match self.0.get::<&str>(key.as_ref()).cloned().or_else(|| { - default - .map(HeaderValue::from_maybe_shared) - .transpose() - .ok() - .flatten() - }) { - Some(value) => HeaderValueBuffer::new(value).into_bytes_ref(py).ok(), - None => None, - } + let value = self + .0 + .get::<&str>(key.as_ref()) + .cloned() + .or_else(|| default.and_then(|d| HeaderValue::from_maybe_shared(d).ok())); + + value.and_then(|v| HeaderValueBuffer::new(v).into_bytes_ref(py).ok()) } /// Returns a view of all values associated with a key. @@ -150,8 +160,7 @@ impl HeaderMap { self.0.is_empty() } - /// Clears the map, removing all key-value pairs. Keeps the allocated memory - /// for reuse. + /// Clears the map, removing all key-value pairs. Keeps the allocated memory for reuse. #[inline] fn clear(&mut self) { self.0.clear(); @@ -191,15 +200,11 @@ impl HeaderMap { inner: self.0.keys().cloned().collect(), } } +} - #[inline] - fn __str__(&self) -> String { - format!("{:?}", self.0) - } - - #[inline] - fn __repr__(&self) -> String { - self.__str__() +impl fmt::Display for HeaderMap { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self.0) } } diff --git a/src/http/mod.rs b/src/http/mod.rs index 2be9eea9..6d97283b 100644 --- a/src/http/mod.rs +++ b/src/http/mod.rs @@ -1,10 +1,9 @@ +pub mod cookie; pub mod header; -mod status; +pub mod status; use pyo3::prelude::*; -pub use self::status::StatusCode; - define_enum!( /// An HTTP version. const, diff --git a/src/http/status.rs b/src/http/status.rs index c753c362..b14036e7 100644 --- a/src/http/status.rs +++ b/src/http/status.rs @@ -1,8 +1,10 @@ +use std::fmt; + use pyo3::prelude::*; /// HTTP status code. -#[pyclass(eq)] #[derive(Clone, Copy, PartialEq, Eq, Hash)] +#[pyclass(eq, str)] pub struct StatusCode(wreq::StatusCode); #[pymethods] @@ -44,19 +46,14 @@ impl StatusCode { } } -#[pymethods] -impl StatusCode { - fn __str__(&self) -> &str { - self.0.as_str() - } - - fn __repr__(&self) -> &str { - self.__str__() - } -} - impl From for StatusCode { fn from(status: wreq::StatusCode) -> Self { Self(status) } } + +impl fmt::Display for StatusCode { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.0) + } +} diff --git a/src/lib.rs b/src/lib.rs index 2694fabd..b9d65e0c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -2,7 +2,7 @@ mod macros; mod buffer; mod client; -mod cookie; + mod emulation; mod error; mod extractor; @@ -22,17 +22,17 @@ use self::{ blocking::{BlockingClient, BlockingResponse, BlockingStreamer, BlockingWebSocket}, delete, get, head, options, patch, post, put, request, trace, websocket, }, - cookie::{Cookie, SameSite}, emulation::{Emulation, EmulationOS, EmulationOption}, error::*, http::{ - Method, StatusCode, Version, + Method, Version, + cookie::{Cookie, Jar, SameSite}, header::{HeaderMap, HeaderMapItemsIter, HeaderMapKeysIter, HeaderMapValuesIter}, + status::StatusCode, }, proxy::Proxy, tls::TlsVersion, }; -use crate::cookie::Jar; #[cfg(all( not(target_env = "msvc"), diff --git a/tests/header_test.py b/tests/header_test.py index d291ed9d..03067359 100644 --- a/tests/header_test.py +++ b/tests/header_test.py @@ -119,3 +119,15 @@ def test_get_with_default(): assert h.get("C") is None assert h.get("C", b"default") == b"default" assert h.get("A", b"default") == b"1" + + +@pytest.mark.flaky(reruns=3, reruns_delay=2) +def test_init_with_dict(): + h = HeaderMap({"A": "1", "B": "2"}) + assert h["A"] == b"1" + assert h["B"] == b"2" + assert len(h) == 2 + assert h.keys_len() == 2 + assert not h.is_empty() + assert h.contains_key("A") + assert h.contains_key("B")