pyo3/conversions/std/
ipaddr.rs

1use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
2
3use crate::exceptions::PyValueError;
4use crate::instance::Bound;
5use crate::sync::GILOnceCell;
6use crate::types::any::PyAnyMethods;
7use crate::types::string::PyStringMethods;
8use crate::types::PyType;
9use crate::{intern, FromPyObject, IntoPy, Py, PyAny, PyObject, PyResult, Python, ToPyObject};
10
11impl FromPyObject<'_> for IpAddr {
12    fn extract_bound(obj: &Bound<'_, PyAny>) -> PyResult<Self> {
13        match obj.getattr(intern!(obj.py(), "packed")) {
14            Ok(packed) => {
15                if let Ok(packed) = packed.extract::<[u8; 4]>() {
16                    Ok(IpAddr::V4(Ipv4Addr::from(packed)))
17                } else if let Ok(packed) = packed.extract::<[u8; 16]>() {
18                    Ok(IpAddr::V6(Ipv6Addr::from(packed)))
19                } else {
20                    Err(PyValueError::new_err("invalid packed length"))
21                }
22            }
23            Err(_) => {
24                // We don't have a .packed attribute, so we try to construct an IP from str().
25                obj.str()?.to_cow()?.parse().map_err(PyValueError::new_err)
26            }
27        }
28    }
29}
30
31impl ToPyObject for Ipv4Addr {
32    fn to_object(&self, py: Python<'_>) -> PyObject {
33        static IPV4_ADDRESS: GILOnceCell<Py<PyType>> = GILOnceCell::new();
34        IPV4_ADDRESS
35            .get_or_try_init_type_ref(py, "ipaddress", "IPv4Address")
36            .expect("failed to load ipaddress.IPv4Address")
37            .call1((u32::from_be_bytes(self.octets()),))
38            .expect("failed to construct ipaddress.IPv4Address")
39            .unbind()
40    }
41}
42
43impl ToPyObject for Ipv6Addr {
44    fn to_object(&self, py: Python<'_>) -> PyObject {
45        static IPV6_ADDRESS: GILOnceCell<Py<PyType>> = GILOnceCell::new();
46        IPV6_ADDRESS
47            .get_or_try_init_type_ref(py, "ipaddress", "IPv6Address")
48            .expect("failed to load ipaddress.IPv6Address")
49            .call1((u128::from_be_bytes(self.octets()),))
50            .expect("failed to construct ipaddress.IPv6Address")
51            .unbind()
52    }
53}
54
55impl ToPyObject for IpAddr {
56    fn to_object(&self, py: Python<'_>) -> PyObject {
57        match self {
58            IpAddr::V4(ip) => ip.to_object(py),
59            IpAddr::V6(ip) => ip.to_object(py),
60        }
61    }
62}
63
64impl IntoPy<PyObject> for IpAddr {
65    fn into_py(self, py: Python<'_>) -> PyObject {
66        self.to_object(py)
67    }
68}
69
70#[cfg(test)]
71mod test_ipaddr {
72    use std::str::FromStr;
73
74    use crate::types::PyString;
75
76    use super::*;
77
78    #[test]
79    fn test_roundtrip() {
80        Python::with_gil(|py| {
81            fn roundtrip(py: Python<'_>, ip: &str) {
82                let ip = IpAddr::from_str(ip).unwrap();
83                let py_cls = if ip.is_ipv4() {
84                    "IPv4Address"
85                } else {
86                    "IPv6Address"
87                };
88
89                let pyobj = ip.into_py(py);
90                let repr = pyobj.bind(py).repr().unwrap();
91                let repr = repr.to_string_lossy();
92                assert_eq!(repr, format!("{}('{}')", py_cls, ip));
93
94                let ip2: IpAddr = pyobj.extract(py).unwrap();
95                assert_eq!(ip, ip2);
96            }
97            roundtrip(py, "127.0.0.1");
98            roundtrip(py, "::1");
99            roundtrip(py, "0.0.0.0");
100        });
101    }
102
103    #[test]
104    fn test_from_pystring() {
105        Python::with_gil(|py| {
106            let py_str = PyString::new_bound(py, "0:0:0:0:0:0:0:1");
107            let ip: IpAddr = py_str.to_object(py).extract(py).unwrap();
108            assert_eq!(ip, IpAddr::from_str("::1").unwrap());
109
110            let py_str = PyString::new_bound(py, "invalid");
111            assert!(py_str.to_object(py).extract::<IpAddr>(py).is_err());
112        });
113    }
114}