radixtarget_rust/
lib.rs

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

radixtarget_rust/
lib.rs
