radixtarget/
lib.rs

1pub mod dns;
2pub mod ip;
3pub mod node;
4pub mod target;
5pub mod utils;
6
7pub use dns::ScopeMode;
8pub use target::RadixTarget;
9
10#[cfg(feature = "py")]
11use pyo3::prelude::*;
12#[cfg(feature = "py")]
13use pyo3::types::PyList;
14#[cfg(feature = "py")]
15use utils::host_size_key;
16
17#[cfg(feature = "py")]
18#[pyclass]
19struct PyRadixTarget {
20    inner: RadixTarget,
21}
22
23#[cfg(feature = "py")]
24#[pymethods]
25impl PyRadixTarget {
26    #[new]
27    #[pyo3(signature = (hosts = None, strict_scope = false, acl_mode = false))]
28    fn new(hosts: Option<Bound<'_, PyList>>, strict_scope: bool, acl_mode: bool) -> PyResult<Self> {
29        let scope_mode = match (strict_scope, acl_mode) {
30            (false, false) => ScopeMode::Normal,
31            (true, false) => ScopeMode::Strict,
32            (false, true) => ScopeMode::Acl,
33            (true, true) => {
34                return Err(pyo3::exceptions::PyValueError::new_err(
35                    "strict_scope and acl_mode are mutually exclusive",
36                ));
37            }
38        };
39        let mut inner =
40            RadixTarget::new(&[], scope_mode).map_err(pyo3::exceptions::PyValueError::new_err)?;
41        if let Some(hosts_list) = hosts {
42            for host in hosts_list.iter() {
43                inner
44                    .insert(&host.extract::<String>()?)
45                    .map_err(pyo3::exceptions::PyValueError::new_err)?;
46            }
47        }
48        Ok(PyRadixTarget { inner })
49    }
50
51    fn insert(&mut self, value: &str) -> PyResult<Option<String>> {
52        self.inner
53            .insert(value)
54            .map_err(pyo3::exceptions::PyValueError::new_err)
55    }
56
57    fn len(&self) -> usize {
58        self.inner.len()
59    }
60
61    fn is_empty(&self) -> bool {
62        self.inner.is_empty()
63    }
64
65    fn contains(&self, value: &str) -> bool {
66        self.inner.contains(value)
67    }
68
69    fn delete(&mut self, value: &str) -> bool {
70        self.inner.delete(value)
71    }
72
73    fn get(&self, value: &str) -> Option<String> {
74        self.inner.get(value)
75    }
76
77    fn prune(&mut self) -> usize {
78        self.inner.prune()
79    }
80
81    fn defrag(&mut self) -> (Vec<String>, Vec<String>) {
82        let (cleaned, new) = self.inner.defrag();
83        (cleaned.into_iter().collect(), new.into_iter().collect())
84    }
85
86    fn __repr__(&self) -> String {
87        format!(
88            "RadixTarget(strict_scope={}, {} hosts)",
89            self.inner.strict_scope(),
90            self.inner.len()
91        )
92    }
93
94    fn __eq__(&self, other: &PyRadixTarget) -> bool {
95        self.inner == other.inner
96    }
97
98    fn __hash__(&self) -> u64 {
99        self.inner.hash()
100    }
101
102    fn contains_target(&self, other: &PyRadixTarget) -> bool {
103        self.inner.contains_target(&other.inner)
104    }
105
106    fn __iter__(slf: PyRef<'_, Self>) -> PyResult<PyRadixTargetIterator> {
107        let hosts: Vec<String> = slf.inner.hosts().iter().cloned().collect();
108        Ok(PyRadixTargetIterator { hosts, index: 0 })
109    }
110
111    fn __bool__(&self) -> bool {
112        !self.inner.is_empty()
113    }
114
115    fn __len__(&self) -> usize {
116        self.inner.len()
117    }
118
119    fn __str__(&self) -> String {
120        let mut hosts: Vec<String> = self.inner.hosts().iter().cloned().collect();
121        hosts.sort();
122
123        if hosts.len() <= 5 {
124            hosts.join(",")
125        } else {
126            format!("{},…", hosts[..5].join(","))
127        }
128    }
129
130    fn copy(&self) -> PyRadixTarget {
131        PyRadixTarget {
132            inner: self.inner.copy(),
133        }
134    }
135}
136
137#[cfg(feature = "py")]
138#[pyclass]
139struct PyRadixTargetIterator {
140    hosts: Vec<String>,
141    index: usize,
142}
143
144#[cfg(feature = "py")]
145#[pymethods]
146impl PyRadixTargetIterator {
147    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
148        slf
149    }
150
151    fn __next__(&mut self) -> Option<String> {
152        if self.index < self.hosts.len() {
153            let result = self.hosts[self.index].clone();
154            self.index += 1;
155            Some(result)
156        } else {
157            None
158        }
159    }
160}
161
162/// PyO3 wrapper for the host_size_key function
163#[cfg(feature = "py")]
164#[pyfunction]
165fn py_host_size_key(host: &Bound<'_, pyo3::PyAny>) -> PyResult<(i64, String)> {
166    // Convert the input to string - this handles both str and ipaddress objects
167    let host_str = host.str()?.to_string();
168    host_size_key(&host_str).map_err(pyo3::exceptions::PyValueError::new_err)
169}
170
171#[cfg(feature = "py")]
172#[pymodule]
173fn _radixtarget(m: &Bound<'_, PyModule>) -> PyResult<()> {
174    m.add_class::<PyRadixTarget>()?;
175    m.add_function(wrap_pyfunction!(py_host_size_key, m)?)?;
176    Ok(())
177}