From 980f05a88efc530375e5cb2803eb8c3aaf4d0caa Mon Sep 17 00:00:00 2001 From: Maximilian Hils Date: Thu, 1 Jun 2023 13:29:23 +0200 Subject: [PATCH] Add `IpAddr` <-> `ipaddress.IPv(4/6)Address` conversion --- newsfragments/3197.added.md | 1 + src/conversions/std/ipaddr.rs | 110 ++++++++++++++++++++++++++++++++++ src/conversions/std/mod.rs | 1 + src/sync.rs | 17 +++++- 4 files changed, 128 insertions(+), 1 deletion(-) create mode 100644 newsfragments/3197.added.md create mode 100755 src/conversions/std/ipaddr.rs diff --git a/newsfragments/3197.added.md b/newsfragments/3197.added.md new file mode 100644 index 00000000..eb8affb1 --- /dev/null +++ b/newsfragments/3197.added.md @@ -0,0 +1 @@ +Add support for converting to and from Python's `ipaddress.IPv4Address`/`ipaddress.IPv6Address` and `std::net::IpAddr`. diff --git a/src/conversions/std/ipaddr.rs b/src/conversions/std/ipaddr.rs new file mode 100755 index 00000000..ca3c8728 --- /dev/null +++ b/src/conversions/std/ipaddr.rs @@ -0,0 +1,110 @@ +use std::net::{IpAddr, Ipv4Addr, Ipv6Addr}; + +use crate::exceptions::PyValueError; +use crate::sync::GILOnceCell; +use crate::types::PyType; +use crate::{intern, FromPyObject, IntoPy, Py, PyAny, PyObject, PyResult, Python, ToPyObject}; + +impl FromPyObject<'_> for IpAddr { + fn extract(obj: &PyAny) -> PyResult { + match obj.getattr(intern!(obj.py(), "packed")) { + Ok(packed) => { + if let Ok(packed) = packed.extract::<[u8; 4]>() { + Ok(IpAddr::V4(Ipv4Addr::from(packed))) + } else if let Ok(packed) = packed.extract::<[u8; 16]>() { + Ok(IpAddr::V6(Ipv6Addr::from(packed))) + } else { + Err(PyValueError::new_err("invalid packed length")) + } + } + Err(_) => { + // We don't have a .packed attribute, so we try to construct an IP from str(). + obj.str()?.to_str()?.parse().map_err(PyValueError::new_err) + } + } + } +} + +impl ToPyObject for Ipv4Addr { + fn to_object(&self, py: Python<'_>) -> PyObject { + static IPV4_ADDRESS: GILOnceCell> = GILOnceCell::new(); + IPV4_ADDRESS + .get_or_try_init_type_ref(py, "ipaddress", "IPv4Address") + .expect("failed to load ipaddress.IPv4Address") + .call1((u32::from_be_bytes(self.octets()),)) + .expect("failed to construct ipaddress.IPv4Address") + .to_object(py) + } +} + +impl ToPyObject for Ipv6Addr { + fn to_object(&self, py: Python<'_>) -> PyObject { + static IPV6_ADDRESS: GILOnceCell> = GILOnceCell::new(); + IPV6_ADDRESS + .get_or_try_init_type_ref(py, "ipaddress", "IPv6Address") + .expect("failed to load ipaddress.IPv6Address") + .call1((u128::from_be_bytes(self.octets()),)) + .expect("failed to construct ipaddress.IPv6Address") + .to_object(py) + } +} + +impl ToPyObject for IpAddr { + fn to_object(&self, py: Python<'_>) -> PyObject { + match self { + IpAddr::V4(ip) => ip.to_object(py), + IpAddr::V6(ip) => ip.to_object(py), + } + } +} + +impl IntoPy for IpAddr { + fn into_py(self, py: Python<'_>) -> PyObject { + self.to_object(py) + } +} + +#[cfg(test)] +mod test_ipaddr { + use std::str::FromStr; + + use crate::types::PyString; + + use super::*; + + #[test] + fn test_roundtrip() { + Python::with_gil(|py| { + fn roundtrip(py: Python<'_>, ip: &str) { + let ip = IpAddr::from_str(ip).unwrap(); + let py_cls = if ip.is_ipv4() { + "IPv4Address" + } else { + "IPv6Address" + }; + + let pyobj = ip.into_py(py); + let repr = pyobj.as_ref(py).repr().unwrap().to_string_lossy(); + assert_eq!(repr, format!("{}('{}')", py_cls, ip)); + + let ip2: IpAddr = pyobj.extract(py).unwrap(); + assert_eq!(ip, ip2); + } + roundtrip(py, "127.0.0.1"); + roundtrip(py, "::1"); + roundtrip(py, "0.0.0.0"); + }); + } + + #[test] + fn test_from_pystring() { + Python::with_gil(|py| { + let py_str = PyString::new(py, "0:0:0:0:0:0:0:1"); + let ip: IpAddr = py_str.to_object(py).extract(py).unwrap(); + assert_eq!(ip, IpAddr::from_str("::1").unwrap()); + + let py_str = PyString::new(py, "invalid"); + assert!(py_str.to_object(py).extract::(py).is_err()); + }); + } +} diff --git a/src/conversions/std/mod.rs b/src/conversions/std/mod.rs index 6021c395..f5e917d0 100644 --- a/src/conversions/std/mod.rs +++ b/src/conversions/std/mod.rs @@ -1,4 +1,5 @@ mod array; +mod ipaddr; mod map; mod num; mod osstr; diff --git a/src/sync.rs b/src/sync.rs index 0f5a5163..3cb4206d 100644 --- a/src/sync.rs +++ b/src/sync.rs @@ -1,5 +1,5 @@ //! Synchronization mechanisms based on the Python GIL. -use crate::{types::PyString, Py, Python}; +use crate::{types::PyString, types::PyType, Py, PyErr, Python}; use std::cell::UnsafeCell; /// Value with concurrent access protected by the GIL. @@ -169,6 +169,21 @@ impl GILOnceCell { } } +impl GILOnceCell> { + /// Get a reference to the contained Python type, initializing it if needed. + /// + /// This is a shorthand method for `get_or_init` which imports the type from Python on init. + pub(crate) fn get_or_try_init_type_ref<'py>( + &'py self, + py: Python<'py>, + module_name: &str, + attr_name: &str, + ) -> Result<&'py PyType, PyErr> { + self.get_or_try_init(py, || py.import(module_name)?.getattr(attr_name)?.extract()) + .map(|ty| ty.as_ref(py)) + } +} + /// Interns `text` as a Python string and stores a reference to it in static storage. /// /// A reference to the same Python string is returned on each invocation.