pyo3/conversions/std/
ipaddr.rs

1use crate::conversion::IntoPyObject;
2use crate::exceptions::PyValueError;
3#[cfg(feature = "experimental-inspect")]
4use crate::inspect::{type_hint_identifier, type_hint_union, PyStaticExpr};
5use crate::sync::PyOnceLock;
6use crate::types::any::PyAnyMethods;
7use crate::types::string::PyStringMethods;
8use crate::types::PyType;
9use crate::{intern, Borrowed, Bound, FromPyObject, Py, PyAny, PyErr, Python};
10use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
11
12impl FromPyObject<'_, '_> for IpAddr {
13    type Error = PyErr;
14
15    #[cfg(feature = "experimental-inspect")]
16    const INPUT_TYPE: PyStaticExpr = type_hint_union!(
17        type_hint_identifier!("ipaddress", "IPv4Address"),
18        type_hint_identifier!("ipaddress", "IPv6Address")
19    );
20
21    fn extract(obj: Borrowed<'_, '_, PyAny>) -> Result<Self, Self::Error> {
22        match obj.getattr(intern!(obj.py(), "packed")) {
23            Ok(packed) => {
24                if let Ok(packed) = packed.extract::<[u8; 4]>() {
25                    Ok(IpAddr::V4(Ipv4Addr::from(packed)))
26                } else if let Ok(packed) = packed.extract::<[u8; 16]>() {
27                    Ok(IpAddr::V6(Ipv6Addr::from(packed)))
28                } else {
29                    Err(PyValueError::new_err("invalid packed length"))
30                }
31            }
32            Err(_) => {
33                // We don't have a .packed attribute, so we try to construct an IP from str().
34                obj.str()?.to_cow()?.parse().map_err(PyValueError::new_err)
35            }
36        }
37    }
38}
39
40impl<'py> IntoPyObject<'py> for Ipv4Addr {
41    type Target = PyAny;
42    type Output = Bound<'py, Self::Target>;
43    type Error = PyErr;
44
45    #[cfg(feature = "experimental-inspect")]
46    const OUTPUT_TYPE: PyStaticExpr = type_hint_identifier!("ipaddress", "IPv4Address");
47
48    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
49        static IPV4_ADDRESS: PyOnceLock<Py<PyType>> = PyOnceLock::new();
50        IPV4_ADDRESS
51            .import(py, "ipaddress", "IPv4Address")?
52            .call1((u32::from_be_bytes(self.octets()),))
53    }
54}
55
56impl<'py> IntoPyObject<'py> for &Ipv4Addr {
57    type Target = PyAny;
58    type Output = Bound<'py, Self::Target>;
59    type Error = PyErr;
60
61    #[cfg(feature = "experimental-inspect")]
62    const OUTPUT_TYPE: PyStaticExpr = Ipv4Addr::OUTPUT_TYPE;
63
64    #[inline]
65    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
66        (*self).into_pyobject(py)
67    }
68}
69
70impl<'py> IntoPyObject<'py> for Ipv6Addr {
71    type Target = PyAny;
72    type Output = Bound<'py, Self::Target>;
73    type Error = PyErr;
74
75    #[cfg(feature = "experimental-inspect")]
76    const OUTPUT_TYPE: PyStaticExpr = type_hint_identifier!("ipaddress", "IPv6Address");
77
78    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
79        static IPV6_ADDRESS: PyOnceLock<Py<PyType>> = PyOnceLock::new();
80        IPV6_ADDRESS
81            .import(py, "ipaddress", "IPv6Address")?
82            .call1((u128::from_be_bytes(self.octets()),))
83    }
84}
85
86impl<'py> IntoPyObject<'py> for &Ipv6Addr {
87    type Target = PyAny;
88    type Output = Bound<'py, Self::Target>;
89    type Error = PyErr;
90
91    #[cfg(feature = "experimental-inspect")]
92    const OUTPUT_TYPE: PyStaticExpr = Ipv6Addr::OUTPUT_TYPE;
93
94    #[inline]
95    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
96        (*self).into_pyobject(py)
97    }
98}
99
100impl<'py> IntoPyObject<'py> for IpAddr {
101    type Target = PyAny;
102    type Output = Bound<'py, Self::Target>;
103    type Error = PyErr;
104
105    #[cfg(feature = "experimental-inspect")]
106    const OUTPUT_TYPE: PyStaticExpr =
107        type_hint_union!(&Ipv4Addr::OUTPUT_TYPE, &Ipv6Addr::OUTPUT_TYPE);
108
109    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
110        match self {
111            IpAddr::V4(ip) => ip.into_pyobject(py),
112            IpAddr::V6(ip) => ip.into_pyobject(py),
113        }
114    }
115}
116
117impl<'py> IntoPyObject<'py> for &IpAddr {
118    type Target = PyAny;
119    type Output = Bound<'py, Self::Target>;
120    type Error = PyErr;
121
122    #[cfg(feature = "experimental-inspect")]
123    const OUTPUT_TYPE: PyStaticExpr = IpAddr::OUTPUT_TYPE;
124
125    #[inline]
126    fn into_pyobject(self, py: Python<'py>) -> Result<Self::Output, Self::Error> {
127        (*self).into_pyobject(py)
128    }
129}
130
131#[cfg(test)]
132mod test_ipaddr {
133    use std::str::FromStr;
134
135    use crate::types::PyString;
136
137    use super::*;
138
139    #[test]
140    fn test_roundtrip() {
141        Python::attach(|py| {
142            fn roundtrip(py: Python<'_>, ip: &str) {
143                let ip = IpAddr::from_str(ip).unwrap();
144                let py_cls = if ip.is_ipv4() {
145                    "IPv4Address"
146                } else {
147                    "IPv6Address"
148                };
149
150                let pyobj = ip.into_pyobject(py).unwrap();
151                let repr = pyobj.repr().unwrap();
152                let repr = repr.to_string_lossy();
153                assert_eq!(repr, format!("{py_cls}('{ip}')"));
154
155                let ip2: IpAddr = pyobj.extract().unwrap();
156                assert_eq!(ip, ip2);
157            }
158            roundtrip(py, "127.0.0.1");
159            roundtrip(py, "::1");
160            roundtrip(py, "0.0.0.0");
161        });
162    }
163
164    #[test]
165    fn test_from_pystring() {
166        Python::attach(|py| {
167            let py_str = PyString::new(py, "0:0:0:0:0:0:0:1");
168            let ip: IpAddr = py_str.extract().unwrap();
169            assert_eq!(ip, IpAddr::from_str("::1").unwrap());
170
171            let py_str = PyString::new(py, "invalid");
172            assert!(py_str.extract::<IpAddr>().is_err());
173        });
174    }
175}