working

src/download/jpmaqsdownload.rs

@@ -1,8 +1,9 @@
 use crate::download::oauth_client::OAuthClient;
 use crate::download::requester::DQRequester;
-use crate::download::timeseries::DQTimeSeriesResponse;
-use crate::download::timeseries::DQTimeseriesRequestArgs;
-use crate::download::timeseries::JPMaQSIndicator;
+use crate::download::helpers::DQTimeSeriesResponse;
+use crate::download::helpers::DQTimeseriesRequestArgs;
+use crate::download::helpers::JPMaQSIndicator;
+// use polars::prelude::*;
 use std::error::Error;
 
 const DEFAULT_JPMAQS_METRICS: [&str; 4] = ["value", "grading", "eop_lag", "mop_lag"];
@@ -59,6 +60,7 @@ impl Default for JPMaQSDownloadGetIndicatorArgs {
     }
 }
 
+/// Struct for downloading data from the JPMaQS data from JPMorgan DataQuery API.
 #[derive(Debug, Clone)]
 pub struct JPMaQSDownload {
     requester: DQRequester,
@@ -72,21 +74,25 @@ impl Default for JPMaQSDownload {
 }
 
 impl JPMaQSDownload {
+    /// Create a new JPMaQSDownload instance with the provided client ID and client secret.
     pub fn new(client_id: String, client_secret: String) -> Self {
         let oauth_client = OAuthClient::new(client_id.clone(), client_secret.clone());
         let requester = DQRequester::new(oauth_client);
         JPMaQSDownload { requester }
     }
 
+    /// Check the connection to the DataQuery API.
     pub fn check_connection(&mut self) -> Result<(), Box<dyn Error>> {
         self.requester.check_connection()
     }
 
+    /// Get the catalogue of tickers available in the JPMaQS data.
     pub fn get_catalogue(&mut self) -> Result<Vec<String>, Box<dyn Error>> {
         let dq_catalogue = self.requester.get_catalogue("JPMAQS", 1000)?;
         Ok(dq_catalogue.all_instruments)
     }
 
+    /// Get the time series data for the provided expressions.
     pub fn get_expressions(
         &mut self,
         expressions: Vec<String>,
@@ -102,7 +108,8 @@ impl JPMaQSDownload {
         Ok(dqts_vec)
     }
 
-    pub fn get_indicators(
+    /// Get the indicators for the provided tickers and metrics.
+    pub fn get_indicators_list(
         &mut self,
         download_args: JPMaQSDownloadGetIndicatorArgs,
     ) -> Result<Vec<JPMaQSIndicator>, Box<dyn Error>> {
@@ -129,4 +136,33 @@ impl JPMaQSDownload {
             Err(e) => Err(e),
         }
     }
+
+    pub fn get_indicators_qdf(
+        &mut self,
+        download_args: JPMaQSDownloadGetIndicatorArgs,
+    ) -> Result<polars::prelude::DataFrame, Box<dyn Error>> {
+        let mut indicators: Vec<JPMaQSIndicator> = self.get_indicators_list(download_args)?;
+        if indicators.is_empty() {
+            return Err("No indicators retrieved".into());
+        }
+        if indicators.len() == 1 {
+            return indicators.pop().unwrap().as_qdf();
+        }
+        assert!(indicators.len() > 1);
+        let mut df_main = indicators.pop().unwrap().as_qdf().unwrap();
+        while !indicators.is_empty() {
+            let df = indicators.pop().unwrap().as_qdf().unwrap();
+            df_main = df_main.vstack(&df).unwrap();
+        }
+        // sort by cid, xcat, real_date in that order
+        let _ = df_main.sort_in_place(
+            [
+                "cid".to_string(),
+                "xcat".to_string(),
+                "real_date".to_string(),
+            ],
+            polars::chunked_array::ops::SortMultipleOptions::default(),
+        );
+        Ok(df_main)
+    }
 }