pyo3/conversions/std/
ipaddr.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};

use crate::exceptions::PyValueError;
use crate::instance::Bound;
use crate::sync::GILOnceCell;
use crate::types::any::PyAnyMethods;
use crate::types::string::PyStringMethods;
use crate::types::PyType;
use crate::{intern, FromPyObject, IntoPy, Py, PyAny, PyObject, PyResult, Python, ToPyObject};

impl FromPyObject<'_> for IpAddr {
    fn extract_bound(obj: &Bound<'_, PyAny>) -> PyResult<Self> {
        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_cow()?.parse().map_err(PyValueError::new_err)
            }
        }
    }
}

impl ToPyObject for Ipv4Addr {
    fn to_object(&self, py: Python<'_>) -> PyObject {
        static IPV4_ADDRESS: GILOnceCell<Py<PyType>> = 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")
            .unbind()
    }
}

impl ToPyObject for Ipv6Addr {
    fn to_object(&self, py: Python<'_>) -> PyObject {
        static IPV6_ADDRESS: GILOnceCell<Py<PyType>> = 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")
            .unbind()
    }
}

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<PyObject> 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.bind(py).repr().unwrap();
                let repr = repr.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_bound(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_bound(py, "invalid");
            assert!(py_str.to_object(py).extract::<IpAddr>(py).is_err());
        });
    }
}