bgpkit_commons/countries/
mod.rs

1//! # Module: countries
2//!
3//! This module provides functionalities related to countries. It fetches country data from the GeoNames database and provides various lookup methods to retrieve country information.
4//!
5//! ## Structures
6//!
7//! ### Country
8//!
9//! This structure represents a country with the following fields:
10//!
11//! - `code`: A 2-letter country code.
12//! - `code3`: A 3-letter country code.
13//! - `name`: The name of the country.
14//! - `capital`: The capital city of the country.
15//! - `continent`: The continent where the country is located.
16//! - `ltd`: The country's top-level domain. This field is optional.
17//! - `neighbors`: A list of neighboring countries represented by their 2-letter country codes.
18//!
19//! ### Countries
20//!
21//! This structure represents a collection of countries. It provides various methods to lookup and retrieve country information.
22use crate::errors::{data_sources, load_methods, modules};
23use crate::{BgpkitCommons, BgpkitCommonsError, LazyLoadable, Result};
24use serde::{Deserialize, Serialize};
25use std::collections::HashMap;
26
27/// Country data structure
28///
29/// Information coming from <https://download.geonames.org/export/dump/countryInfo.txt>
30#[derive(Debug, Clone, Serialize, Deserialize)]
31pub struct Country {
32    /// 2-letter country code
33    pub code: String,
34    /// 3-letter country code
35    pub code3: String,
36    /// Country name
37    pub name: String,
38    /// Capital city
39    pub capital: String,
40    /// Continent
41    pub continent: String,
42    /// Country's top-level domain
43    pub ltd: Option<String>,
44    /// Neighboring countries in 2-letter country code
45    pub neighbors: Vec<String>,
46}
47
48#[derive(Debug, Clone, Serialize, Deserialize)]
49pub struct Countries {
50    countries: HashMap<String, Country>,
51}
52
53const DATA_URL: &str = "https://download.geonames.org/export/dump/countryInfo.txt";
54
55impl Countries {
56    pub fn new() -> Result<Self> {
57        let mut countries: Vec<Country> = vec![];
58        for line in oneio::read_lines(DATA_URL)? {
59            let text = line.ok().ok_or_else(|| {
60                BgpkitCommonsError::data_source_error(data_sources::GEONAMES, "error reading line")
61            })?;
62            if text.trim() == "" || text.starts_with('#') {
63                continue;
64            }
65            let splits: Vec<&str> = text.split('\t').collect();
66            if splits.len() != 19 {
67                return Err(BgpkitCommonsError::invalid_format(
68                    "countries data",
69                    text.as_str(),
70                    "row missing fields",
71                ));
72            }
73            let code = splits[0].to_string();
74            let code3 = splits[1].to_string();
75            let name = splits[4].to_string();
76            let capital = splits[5].to_string();
77            let continent = splits[8].to_string();
78            let ltd = match splits[9] {
79                "" => None,
80                d => Some(d.to_string()),
81            };
82            let neighbors = splits[17]
83                .split(',')
84                .map(|x| x.to_string())
85                .collect::<Vec<String>>();
86            countries.push(Country {
87                code,
88                code3,
89                name,
90                capital,
91                continent,
92                ltd,
93                neighbors,
94            })
95        }
96
97        let mut countries_map: HashMap<String, Country> = HashMap::new();
98        for country in countries {
99            countries_map.insert(country.code.clone(), country);
100        }
101        Ok(Countries {
102            countries: countries_map,
103        })
104    }
105
106    /// Lookup country by 2-letter country code
107    pub fn lookup_by_code(&self, code: &str) -> Option<Country> {
108        self.countries.get(code).cloned()
109    }
110
111    /// Lookup country by country name
112    ///
113    /// This function is case-insensitive and search for countries with name that contains the given name string
114    pub fn lookup_by_name(&self, name: &str) -> Vec<Country> {
115        let lower_name = name.to_lowercase();
116        let mut countries: Vec<Country> = vec![];
117        for country in self.countries.values() {
118            if country.name.to_lowercase().contains(&lower_name) {
119                countries.push(country.clone());
120            }
121        }
122        countries
123    }
124
125    /// Get all countries
126    pub fn all_countries(&self) -> Vec<Country> {
127        self.countries.values().cloned().collect()
128    }
129}
130
131impl LazyLoadable for Countries {
132    fn reload(&mut self) -> Result<()> {
133        *self = Countries::new().map_err(|e| {
134            BgpkitCommonsError::data_source_error(data_sources::GEONAMES, e.to_string())
135        })?;
136        Ok(())
137    }
138
139    fn is_loaded(&self) -> bool {
140        !self.countries.is_empty()
141    }
142
143    fn loading_status(&self) -> &'static str {
144        if self.is_loaded() {
145            "Countries data loaded"
146        } else {
147            "Countries data not loaded"
148        }
149    }
150}
151
152impl BgpkitCommons {
153    pub fn country_all(&self) -> Result<Vec<Country>> {
154        if self.countries.is_none() {
155            return Err(BgpkitCommonsError::module_not_loaded(
156                modules::COUNTRIES,
157                load_methods::LOAD_COUNTRIES,
158            ));
159        }
160
161        Ok(self.countries.as_ref().unwrap().all_countries())
162    }
163
164    pub fn country_by_code(&self, code: &str) -> Result<Option<Country>> {
165        if self.countries.is_none() {
166            return Err(BgpkitCommonsError::module_not_loaded(
167                modules::COUNTRIES,
168                load_methods::LOAD_COUNTRIES,
169            ));
170        }
171        Ok(self.countries.as_ref().unwrap().lookup_by_code(code))
172    }
173
174    pub fn country_by_name(&self, name: &str) -> Result<Vec<Country>> {
175        if self.countries.is_none() {
176            return Err(BgpkitCommonsError::module_not_loaded(
177                modules::COUNTRIES,
178                load_methods::LOAD_COUNTRIES,
179            ));
180        }
181        Ok(self.countries.as_ref().unwrap().lookup_by_name(name))
182    }
183
184    pub fn country_by_code3(&self, code: &str) -> Result<Option<Country>> {
185        if self.countries.is_none() {
186            return Err(BgpkitCommonsError::module_not_loaded(
187                modules::COUNTRIES,
188                load_methods::LOAD_COUNTRIES,
189            ));
190        }
191        Ok(self.countries.as_ref().unwrap().lookup_by_code(code))
192    }
193}