yuzu_data/
csv_io.rs

1//! Per-symbol gzip CSV I/O. Files hold full adjusted OHLCV
2//! (`day,adj_open,adj_high,adj_low,adj_close,volume`, oldest-first); `parse_series`
3//! extracts one chosen [`Field`] column into `(day, value)` rows.
4
5use crate::error::DataError;
6use flate2::read::GzDecoder;
7use flate2::write::GzEncoder;
8use flate2::Compression;
9use std::io::{Read, Write};
10
11/// One row of adjusted OHLCV for a trading day.
12#[derive(Debug, Clone, PartialEq)]
13pub struct OhlcvRow {
14    pub day: i32,
15    pub adj_open: f64,
16    pub adj_high: f64,
17    pub adj_low: f64,
18    pub adj_close: f64,
19    pub volume: f64,
20}
21
22/// Which OHLCV column to pull into a Panel.
23#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24pub enum Field {
25    AdjOpen,
26    AdjHigh,
27    AdjLow,
28    AdjClose,
29    Volume,
30}
31
32impl Field {
33    /// Column index in the CSV (day is column 0).
34    fn col(self) -> usize {
35        match self {
36            Field::AdjOpen => 1,
37            Field::AdjHigh => 2,
38            Field::AdjLow => 3,
39            Field::AdjClose => 4,
40            Field::Volume => 5,
41        }
42    }
43}
44
45fn date_to_i32(s: &str) -> Result<i32, DataError> {
46    s.replace('-', "")
47        .parse()
48        .map_err(|_| DataError::Parse(format!("bad date '{s}'")))
49}
50
51fn i32_to_date(d: i32) -> String {
52    format!("{:04}-{:02}-{:02}", d / 10000, d / 100 % 100, d % 100)
53}
54
55/// gunzip + parse the OHLCV CSV (header optional, oldest-first), extracting
56/// `(YYYYMMDD, <field>)` for the chosen field.
57pub fn parse_series(gz: &[u8], field: Field) -> Result<Vec<(i32, f64)>, DataError> {
58    let mut text = String::new();
59    GzDecoder::new(gz)
60        .read_to_string(&mut text)
61        .map_err(|e| DataError::Io(e.to_string()))?;
62    let col = field.col();
63    let mut out = Vec::new();
64    for line in text.lines() {
65        let line = line.trim();
66        if line.is_empty() || line.starts_with("day") {
67            continue;
68        }
69        let cells: Vec<&str> = line.split(',').collect();
70        if cells.len() <= col {
71            return Err(DataError::Parse(format!(
72                "row has {} cells, need > {col}",
73                cells.len()
74            )));
75        }
76        let day = date_to_i32(cells[0].trim())?;
77        let val: f64 = cells[col]
78            .trim()
79            .parse()
80            .map_err(|_| DataError::Parse(format!("bad value '{}'", cells[col])))?;
81        out.push((day, val));
82    }
83    Ok(out)
84}
85
86/// Serialize OHLCV rows (oldest-first) to gzip CSV with the standard header.
87pub fn write_series(rows: &[OhlcvRow]) -> Result<Vec<u8>, DataError> {
88    let mut buf = String::from("day,adj_open,adj_high,adj_low,adj_close,volume\n");
89    for r in rows {
90        buf.push_str(&format!(
91            "{},{},{},{},{},{}\n",
92            i32_to_date(r.day),
93            r.adj_open,
94            r.adj_high,
95            r.adj_low,
96            r.adj_close,
97            r.volume
98        ));
99    }
100    let mut enc = GzEncoder::new(Vec::new(), Compression::default());
101    enc.write_all(buf.as_bytes())
102        .map_err(|e| DataError::Io(e.to_string()))?;
103    enc.finish().map_err(|e| DataError::Io(e.to_string()))
104}
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109
110    fn row(day: i32, c: f64) -> OhlcvRow {
111        OhlcvRow {
112            day,
113            adj_open: c - 0.5,
114            adj_high: c + 1.0,
115            adj_low: c - 1.0,
116            adj_close: c,
117            volume: 1000.0,
118        }
119    }
120
121    #[test]
122    fn round_trips_and_selects_fields() {
123        let rows = vec![row(20240102, 10.5), row(20240103, 11.25)];
124        let gz = write_series(&rows).unwrap();
125        assert_eq!(
126            parse_series(&gz, Field::AdjClose).unwrap(),
127            vec![(20240102, 10.5), (20240103, 11.25)]
128        );
129        assert_eq!(
130            parse_series(&gz, Field::AdjHigh).unwrap(),
131            vec![(20240102, 11.5), (20240103, 12.25)]
132        );
133        assert_eq!(
134            parse_series(&gz, Field::AdjLow).unwrap(),
135            vec![(20240102, 9.5), (20240103, 10.25)]
136        );
137        assert_eq!(
138            parse_series(&gz, Field::AdjOpen).unwrap(),
139            vec![(20240102, 10.0), (20240103, 10.75)]
140        );
141        assert_eq!(
142            parse_series(&gz, Field::Volume).unwrap(),
143            vec![(20240102, 1000.0), (20240103, 1000.0)]
144        );
145    }
146
147    fn gz(text: &str) -> Vec<u8> {
148        let mut e = GzEncoder::new(Vec::new(), Compression::default());
149        e.write_all(text.as_bytes()).unwrap();
150        e.finish().unwrap()
151    }
152
153    #[test]
154    fn parse_errors_on_non_gzip_bad_value_bad_date_and_short_row() {
155        let hdr = "day,adj_open,adj_high,adj_low,adj_close,volume\n";
156        assert!(parse_series(b"not gzip", Field::AdjClose).is_err());
157        assert!(parse_series(
158            &gz(&format!("{hdr}2024-01-02,1,2,3,bad,5\n")),
159            Field::AdjClose
160        )
161        .is_err());
162        assert!(parse_series(&gz(&format!("{hdr}bad-date,1,2,3,4,5\n")), Field::AdjClose).is_err());
163        // row missing the volume column, but Volume requested
164        assert!(parse_series(&gz(&format!("{hdr}2024-01-02,1,2,3,4\n")), Field::Volume).is_err());
165    }
166}
yuzu_data/csv_io.rs

yuzu_data/
csv_io.rs
