Skip to main content

pomelo_data/
industry.rs

1//! Build the `symbol -> sector` map the yuzu-core neutralization ops need,
2//! from the `tracked/*.csv.gz` snapshot (`symbol,sector,market_cap`).
3
4use std::collections::HashMap;
5
6use crate::error::DataError;
7use crate::format::read_csv_text;
8use crate::source::ObjectSource;
9
10/// Load `tracked/universe.csv.gz` (or `.csv`) from `src` into the
11/// `symbol → sector` map the engine's industry ops consume. `Ok(None)` when
12/// the snapshot is absent or has no usable rows — the runner then leaves the
13/// map empty and industry ops treat every symbol as unmapped.
14pub 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
26/// Parse a `symbol,sector,...` CSV into `symbol -> sector`. A leading header row
27/// (`symbol,...`) is skipped; rows with an empty sector are dropped.
28pub 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")); // empty sector dropped
59        assert!(!m.contains_key("symbol")); // header skipped
60        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        // Absent tree → no map, no error.
73        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}