stock_data/
lib.rs

1use chrono::{NaiveDate, TimeZone, Utc};
2use reqwest::Error as ReqwestError;
3use reqwest::header::{HeaderMap, USER_AGENT};
4use std::io::Error as IoError;
5use tokio::fs::File;
6use tokio::io::AsyncWriteExt;
7
8/// Writes the provided bytes to the specified file path.
9///
10/// # Arguments
11///
12/// * `bytes` - The bytes to be written to the file.
13/// * `path` - The path of the file to write the bytes to.
14///
15/// # Errors
16///
17/// Returns an `IoError` if there was an error writing to the file.
18pub async fn write_stock_data(bytes: &[u8], path: &str) -> Result<(), IoError> {
19    let mut file = File::create(path).await?;
20    file.write_all(&bytes).await?;
21    Ok(())
22}
23
24/// Downloads stock data from the specified URL.
25///
26/// # Arguments
27///
28/// * `url` - The URL to download the stock data from.
29///
30/// # Errors
31///
32/// Returns a `ReqwestError` if there was an error downloading the data.
33pub async fn download_stock_data(url: &str) -> Result<Vec<u8>, ReqwestError> {
34    let mut headers = HeaderMap::new();
35    headers.insert(USER_AGENT, "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/89.0.4389.82 Safari/537.36".parse().unwrap());
36
37    let client = reqwest::Client::new();
38    let response = client.get(url)
39        .headers(headers)
40        .send()
41        .await?;
42
43    if response.status().is_success() {
44        let bytes = response.bytes().await?;
45        Ok(bytes.to_vec())
46    } else {
47        Err(response.error_for_status().unwrap_err())
48    }
49}
50
51/// Builds a Yahoo Finance URL for downloading stock data based on the provided dates.
52///
53/// # Arguments
54///
55/// * `symbol` - The stock symbol.
56/// * `date1` - The start date.
57/// * `date2` - The end date.
58/// * `interval` - The interval between data points (e.g., "1d" for daily).
59/// * `include_adjusted_close` - Whether to include the adjusted close price.
60///
61/// # Returns
62///
63/// The generated Yahoo Finance URL.
64pub fn build_yahoo_finance_url_from_dates(
65    symbol: &str,
66    date1: NaiveDate,
67    date2: NaiveDate,
68    interval: &str,
69    include_adjusted_close: bool,
70) -> String {
71    let msg = "Failed to generate timestamp from date";
72    let datetime1 = date1.and_hms_opt(0, 0, 0).expect(msg);
73    let datetime2 = date2.and_hms_opt(0, 0, 0).expect(msg);
74
75    let period1 = Utc.from_utc_datetime(&datetime1).timestamp() as u64;
76    let period2 = Utc.from_utc_datetime(&datetime2).timestamp() as u64;
77    build_yahoo_finance_url(symbol, period1, period2, interval, include_adjusted_close)
78}
79
80/// Builds a Yahoo Finance URL for downloading stock data based on the provided timestamps.
81///
82/// # Arguments
83///
84/// * `symbol` - The stock symbol.
85/// * `date1` - The start timestamp.
86/// * `date2` - The end timestamp.
87/// * `interval` - The interval between data points (e.g., "1d" for daily).
88/// * `include_adjusted_close` - Whether to include the adjusted close price.
89///
90/// # Returns
91///
92/// The generated Yahoo Finance URL.
93pub fn build_yahoo_finance_url(
94    symbol: &str,
95    period1: u64,
96    period2: u64,
97    interval: &str,
98    include_adjusted_close: bool,
99) -> String {
100    format!(
101        "https://query1.finance.yahoo.com/v7/finance/download/{}?period1={}&period2={}&interval={}&events=history&includeAdjustedClose={}",
102        symbol,
103        period1,
104        period2,
105        interval,
106        include_adjusted_close
107    )
108}
109
110#[cfg(test)]
111mod tests {
112    use super::*;
113    use chrono::NaiveDate;
114    use std::fs::File;
115    use std::io::Read;
116
117    #[test]
118    fn test_build_yahoo_finance_url_from_dates() {
119        let symbol = "AMZN";
120        let date1 = NaiveDate::from_ymd_opt(2021, 1, 1).unwrap();
121        let date2 = NaiveDate::from_ymd_opt(2022, 1, 1).unwrap();
122        let interval = "1d";
123        let include_adjusted_close = true;
124
125        let url = build_yahoo_finance_url_from_dates(symbol, date1, date2, interval, include_adjusted_close);
126
127        assert_eq!(
128            url,
129            "https://query1.finance.yahoo.com/v7/finance/download/AMZN?period1=1609459200&period2=1640995200&interval=1d&events=history&includeAdjustedClose=true"
130        );
131    }
132
133    #[test]
134    fn test_build_yahoo_finance_url() {
135        let symbol = "AMZN";
136        let period1 = 1609459200;
137        let period2 = 1640995200;
138        let interval = "1d";
139        let include_adjusted_close = true;
140
141        let url = build_yahoo_finance_url(symbol, period1, period2, interval, include_adjusted_close);
142
143        assert_eq!(
144            url,
145            "https://query1.finance.yahoo.com/v7/finance/download/AMZN?period1=1609459200&period2=1640995200&interval=1d&events=history&includeAdjustedClose=true"
146        );
147    }
148
149    #[tokio::test]
150    async fn test_download_and_write_stock_data() {
151        let symbol = "AMZN";
152        let date1 = NaiveDate::from_ymd_opt(2021, 1, 1).unwrap();
153        let date2 = NaiveDate::from_ymd_opt(2022, 1, 1).unwrap();
154        let interval = "1d";
155        let include_adjusted_close = true;
156
157        let url = build_yahoo_finance_url_from_dates(symbol, date1, date2, interval, include_adjusted_close);
158
159        // Test downloading stock data
160        let data = download_stock_data(&url).await.unwrap();
161        assert!(!data.is_empty());
162
163        // Test writing stock data to a file
164        let path = "test_data.csv";
165        write_stock_data(&data, path).await.unwrap();
166
167        // Read the written file and verify its content
168        let mut file = File::open(path).unwrap();
169        let mut content = String::new();
170        file.read_to_string(&mut content).unwrap();
171        assert!(!content.is_empty());
172
173        // Clean up the test file
174        std::fs::remove_file(path).unwrap();
175    }
176}