1use std::collections::HashMap;
5
6use crate::error::DataError;
7use crate::format::read_csv_text;
8use crate::source::ObjectSource;
9
10pub fn load_industry_map<S: ObjectSource>(
15 src: &S,
16) -> Result<Option<HashMap<String, String>>, DataError> {
17 for ext in [".csv.gz", ".csv"] {
18 if let Some(bytes) = src.get(&format!("tracked/universe{ext}"))? {
19 let map = parse_industry_csv(&read_csv_text(&bytes)?);
20 return Ok((!map.is_empty()).then_some(map));
21 }
22 }
23 Ok(None)
24}
25
26pub fn parse_industry_csv(csv: &str) -> HashMap<String, String> {
29 let mut out = HashMap::new();
30 for line in csv.lines() {
31 let line = line.trim();
32 if line.is_empty() {
33 continue;
34 }
35 let mut fields = line.split(',');
36 let (Some(sym), Some(sector)) = (fields.next(), fields.next()) else {
37 continue;
38 };
39 let (sym, sector) = (sym.trim(), sector.trim());
40 if sym.is_empty() || sector.is_empty() || sym == "symbol" {
41 continue;
42 }
43 out.insert(sym.to_string(), sector.to_string());
44 }
45 out
46}
47
48#[cfg(test)]
49mod tests {
50 use super::*;
51
52 #[test]
53 fn parses_symbol_sector_skipping_header_and_blanks() {
54 let csv = "symbol,sector,market_cap\nNVDA,Technology,5103000000000\nXOM,Energy,470000000000\nUNKNOWN,,1000\n";
55 let m = parse_industry_csv(csv);
56 assert_eq!(m.get("NVDA").map(String::as_str), Some("Technology"));
57 assert_eq!(m.get("XOM").map(String::as_str), Some("Energy"));
58 assert!(!m.contains_key("UNKNOWN")); assert!(!m.contains_key("symbol")); assert_eq!(m.len(), 2);
61 }
62
63 #[test]
64 fn loads_the_tracked_snapshot_when_present() {
65 use flate2::write::GzEncoder;
66 use flate2::Compression;
67 use std::io::Write;
68
69 let dir = std::env::temp_dir().join("pomelo_data_industry_map");
70 let _ = std::fs::remove_dir_all(&dir);
71 let src = crate::LocalSource::new(&dir);
72 assert!(load_industry_map(&src).unwrap().is_none());
74
75 std::fs::create_dir_all(dir.join("tracked")).unwrap();
76 let mut enc = GzEncoder::new(Vec::new(), Compression::default());
77 enc.write_all(b"symbol,sector,market_cap\nAAA,Tech,1\nBBB,Energy,2\n")
78 .unwrap();
79 std::fs::write(dir.join("tracked/universe.csv.gz"), enc.finish().unwrap()).unwrap();
80 let map = load_industry_map(&src).unwrap().unwrap();
81 assert_eq!(map.get("AAA").map(String::as_str), Some("Tech"));
82 assert_eq!(map.len(), 2);
83 }
84}