Skip to main content

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