geodb_rs/
lib.rs

1#![allow(clippy::useless_conversion)]
2
3use geodb_core::{
4    CityView, CountryView, DefaultGeoDb, GeoDb, PhoneCodeSearch, SmartItem, StandardBackend,
5    StateView,
6};
7use pyo3::exceptions::PyRuntimeError;
8use pyo3::prelude::*;
9use pyo3::types::PyModule;
10use serde::Serialize;
11use std::path::{Path, PathBuf};
12
13// Provide a single helper to convert geodb-core Results into PyResult
14// while avoiding per-call mapping. Due to Rust orphan rules we cannot
15// implement `From<GeoError> for PyErr` in this crate, so we add an
16// extension trait and use `.into_py()?` at call sites.
17trait IntoPyResult<T> {
18    fn into_py(self) -> PyResult<T>;
19}
20
21impl<T> IntoPyResult<T> for geodb_core::Result<T> {
22    fn into_py(self) -> PyResult<T> {
23        self.map_err(|e| PyRuntimeError::new_err(e.to_string()))
24    }
25}
26
27#[pyclass]
28pub struct PyGeoDb {
29    inner: DefaultGeoDb,
30}
31
32fn to_py<'py, T: Serialize + ?Sized>(py: Python<'py>, value: &T) -> PyResult<Bound<'py, PyAny>> {
33    // Serialize to JSON string, then parse in Python via json.loads to get native dict/list
34    let s = serde_json::to_string(value)
35        .map_err(|e| PyErr::new::<PyRuntimeError, _>(format!("serde error: {e}")))?;
36    let json_mod = PyModule::import_bound(py, "json")?;
37    let loads = json_mod.getattr("loads")?;
38    let obj = loads.call1((s,))?;
39    Ok(obj)
40}
41
42/// Find the bundled data file in the Python package
43fn find_bundled_data() -> PyResult<PathBuf> {
44    Python::with_gil(|py| {
45        let geodb_module = py.import_bound("geodb_rs")?;
46        let module_path = geodb_module.getattr("__file__")?.extract::<String>()?;
47
48        let module_dir = Path::new(&module_path).parent().ok_or_else(|| {
49            PyErr::new::<PyRuntimeError, _>("Could not determine module directory")
50        })?;
51
52        // Try multiple possible locations for the data file
53        let possible_paths = [
54            // Location when installed via pip (maturin's data directory)
55            module_dir.join("../geodb_rs_data/countries+states+cities.json.gz"),
56            // Location if data is in the package itself
57            module_dir.join("data/countries+states+cities.json.gz"),
58            // Fallback to sibling directory
59            module_dir.join("geodb_rs_data/countries+states+cities.json.gz"),
60        ];
61
62        for path in &possible_paths {
63            if path.exists() {
64                return path.canonicalize().map_err(|e| {
65                    PyErr::new::<PyRuntimeError, _>(format!("Failed to canonicalize path: {e}"))
66                });
67            }
68        }
69
70        Err(PyErr::new::<pyo3::exceptions::PyFileNotFoundError, _>(
71            format!(
72                "Data file not found. Searched in:\n{}",
73                possible_paths
74                    .iter()
75                    .map(|p| format!("  - {}", p.display()))
76                    .collect::<Vec<_>>()
77                    .join("\n")
78            ),
79        ))
80    })
81}
82
83#[pymethods]
84impl PyGeoDb {
85    #[staticmethod]
86    pub fn load(path: &str) -> PyResult<Self> {
87        let db = GeoDb::<StandardBackend>::load_from_path(path, None).into_py()?;
88        Ok(Self { inner: db })
89    }
90
91    #[staticmethod]
92    pub fn load_default() -> PyResult<Self> {
93        // Try to find bundled data first, fall back to geodb-core's default
94        match find_bundled_data() {
95            Ok(path) => {
96                let path_str = path
97                    .to_str()
98                    .ok_or_else(|| PyRuntimeError::new_err("Invalid path encoding"))?;
99                let db = GeoDb::<StandardBackend>::load_from_path(path_str, None).into_py()?;
100                Ok(Self { inner: db })
101            }
102            Err(_) => {
103                // Fall back to geodb-core's default location
104                let db = GeoDb::<StandardBackend>::load().into_py()?;
105                Ok(Self { inner: db })
106            }
107        }
108    }
109
110    #[staticmethod]
111    pub fn load_filtered(iso2_list: Vec<String>) -> PyResult<Self> {
112        let tmp: Vec<String> = iso2_list
113            .into_iter()
114            .map(|s| s.trim().to_string())
115            .collect();
116        let refs: Vec<&str> = tmp.iter().map(String::as_str).collect();
117        let db = GeoDb::<StandardBackend>::load_filtered_by_iso2(&refs).into_py()?;
118        Ok(Self { inner: db })
119    }
120
121    pub fn stats(&self) -> PyResult<(usize, usize, usize)> {
122        let s = self.inner.stats();
123        Ok((s.countries, s.states, s.cities))
124    }
125
126    /// Return a list of all countries as dicts
127    pub fn countries<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyAny>> {
128        let items: Vec<_> = self.inner.countries().iter().map(CountryView).collect();
129        to_py(py, &items)
130    }
131
132    /// Find a country by ISO2/ISO3/code and return as dict (or None)
133    pub fn find_country<'py>(
134        &self,
135        py: Python<'py>,
136        code: &str,
137    ) -> PyResult<Option<Bound<'py, PyAny>>> {
138        if let Some(c) = self.inner.find_country_by_code(code) {
139            let v = to_py(py, &CountryView(c))?;
140            Ok(Some(v))
141        } else {
142            Ok(None)
143        }
144    }
145
146    /// List all states for a given country ISO2 as dicts
147    pub fn states_in_country<'py>(
148        &self,
149        py: Python<'py>,
150        iso2: &str,
151    ) -> PyResult<Option<Bound<'py, PyAny>>> {
152        if let Some(country) = self.inner.find_country_by_iso2(iso2) {
153            let items: Vec<_> = country
154                .states()
155                .iter()
156                .map(|s| StateView { country, state: s })
157                .collect();
158            let obj = to_py(py, &items)?;
159            Ok(Some(obj))
160        } else {
161            Ok(None)
162        }
163    }
164
165    /// Find countries by phone code (e.g. "+49", "1") and return list of dicts
166    pub fn search_countries_by_phone<'py>(
167        &self,
168        py: Python<'py>,
169        phone: &str,
170    ) -> PyResult<Bound<'py, PyAny>> {
171        let code = phone.trim().trim_start_matches('+');
172        let items: Vec<_> = self
173            .inner
174            .find_countries_by_phone_code(code)
175            .iter()
176            .map(|c| CountryView(*c))
177            .collect();
178        to_py(py, &items)
179    }
180
181    /// Find states containing a substring (ASCII, case-insensitive). Returns list of dicts
182    pub fn find_states_by_substring<'py>(
183        &self,
184        py: Python<'py>,
185        substr: &str,
186    ) -> PyResult<Bound<'py, PyAny>> {
187        let items: Vec<_> = self
188            .inner
189            .find_states_by_substring(substr)
190            .into_iter()
191            .map(|(state, country)| StateView { country, state })
192            .collect();
193        to_py(py, &items)
194    }
195
196    /// Find cities containing a substring (ASCII, case-insensitive). Returns list of dicts
197    pub fn find_cities_by_substring<'py>(
198        &self,
199        py: Python<'py>,
200        substr: &str,
201    ) -> PyResult<Bound<'py, PyAny>> {
202        let items: Vec<_> = self
203            .inner
204            .find_cities_by_substring(substr)
205            .into_iter()
206            .map(|(city, state, country)| CityView {
207                country,
208                state,
209                city,
210            })
211            .collect();
212        to_py(py, &items)
213    }
214
215    /// Smart search across countries, states, cities, and phone codes. Returns list of dicts
216    pub fn smart_search<'py>(&self, py: Python<'py>, query: &str) -> PyResult<Bound<'py, PyAny>> {
217        let hits = self.inner.smart_search(query);
218        // Map to a homogeneous list by emitting the view of the matched entity
219        let mut out: Vec<serde_json::Value> = Vec::with_capacity(hits.len());
220        for hit in hits {
221            let v = match hit.item {
222                SmartItem::Country(c) => serde_json::to_value(CountryView(c))
223                    .map_err(|e| PyRuntimeError::new_err(format!("serde error: {e}")))?,
224                SmartItem::State { country, state } => {
225                    serde_json::to_value(&StateView { country, state })
226                        .map_err(|e| PyRuntimeError::new_err(format!("serde error: {e}")))?
227                }
228                SmartItem::City {
229                    country,
230                    state,
231                    city,
232                } => serde_json::to_value(&CityView {
233                    country,
234                    state,
235                    city,
236                })
237                .map_err(|e| PyRuntimeError::new_err(format!("serde error: {e}")))?,
238            };
239            out.push(v);
240        }
241        to_py(py, &out)
242    }
243}
244
245/// Python module entry point
246/// IMPORTANT: The function name must match the module-name in pyproject.toml
247#[pymodule]
248fn geodb_rs(_py: Python, m: &Bound<PyModule>) -> PyResult<()> {
249    m.add_class::<PyGeoDb>()?;
250    Ok(())
251}
geodb_rs/lib.rs

geodb_rs/
lib.rs
