From 213ea72b2083b6ce6f0acda287ac1eb29602d9bd Mon Sep 17 00:00:00 2001
From: Palash Tyagi <23239946+Magnus167@users.noreply.github.com>
Date: Mon, 18 Nov 2024 01:02:51 +0000
Subject: [PATCH] add Python wrapper for downloading JPMaQS indicators as
 DataFrame

---
 src/download/jpmaqsdownload.rs | 38 ++++++++++++++++++++++++++++++
 src/py/download.rs             | 43 ++++++++++++++++++++++++++++++++++
 2 files changed, 81 insertions(+)
 create mode 100644 src/py/download.rs

diff --git a/src/download/jpmaqsdownload.rs b/src/download/jpmaqsdownload.rs
index 0fd8b47..a485389 100644
--- a/src/download/jpmaqsdownload.rs
+++ b/src/download/jpmaqsdownload.rs
@@ -368,3 +368,41 @@ fn save_indicators_list_as_csvs(
     }
     Ok(())
 }
+
+/// Function to download the JPMaQS tickers as a DataFrame. Solely exists to allow for a simpler
+/// Python wrapper for downloading JPMaQS data.
+pub fn download_jpmaqs_indicators_as_df(
+    client_id: String,
+    client_secret: String,
+    tickers: Vec<String>,
+    metrics: Option<Vec<String>>,
+    start_date: Option<String>,
+    end_date: Option<String>,
+) -> Result<polars::prelude::DataFrame, Box<dyn Error>> {
+    let metrics = metrics.unwrap_or_else(|| {
+        DEFAULT_JPMAQS_METRICS
+            .iter()
+            .map(|s| s.to_string())
+            .collect()
+    });
+
+    // if metrics== "ALL", then use the default metrics
+    let metrics = if metrics.len() == 1 && metrics[0] == "ALL" {
+        DEFAULT_JPMAQS_METRICS
+            .iter()
+            .map(|s| s.to_string())
+            .collect()
+    } else {
+        metrics
+    };
+
+    let mut jpamqs_download = JPMaQSDownload::new(client_id, client_secret);
+    let download_args = JPMaQSDownloadGetIndicatorArgs {
+        tickers: tickers,
+        metrics: metrics,
+        start_date: start_date.unwrap_or("1990-01-01".to_string()),
+        end_date: end_date.unwrap_or("TODAY+2D".to_string()),
+        ..Default::default()
+    };
+    jpamqs_download.get_indicators_qdf(download_args)
+}
diff --git a/src/py/download.rs b/src/py/download.rs
new file mode 100644
index 0000000..5600fff
--- /dev/null
+++ b/src/py/download.rs
@@ -0,0 +1,43 @@
+use pyo3::prelude::*;
+use pyo3_polars::PyDataFrame;
+
+/// Python wrapper for [`crate::utils::qdf`] module.
+#[allow(deprecated)]
+#[pymodule]
+pub fn download(_py: Python, m: &PyModule) -> PyResult<()> {
+    m.add_function(wrap_pyfunction!(download_jpmaqs_indicators_as_df, m)?)?;
+    Ok(())
+}
+
+/// Python wrapper for downloading JPMAQS indicators as a DataFrame.
+/// See [`crate::download::jpmaqsdownload::download_jpmaqs_indicators_as_df`] for full documentation.
+/// Arguments:
+/// - `client_id`: The client ID.
+/// - `client_secret`: The client secret.
+/// - `tickers`: The tickers.
+/// - `start_date`: The start date.
+/// - `end_date`: The end date.
+/// 
+/// Returns:
+/// - A standardized Quantamental DataFrame.
+#[pyfunction]
+pub fn download_jpmaqs_indicators_as_df(
+    client_id: String,
+    client_secret: String,
+    tickers: Vec<String>,
+    metrics: Option<Vec<String>>,
+    start_date: Option<String>,
+    end_date: Option<String>,
+) -> PyResult<PyDataFrame> {
+    Ok(PyDataFrame(
+        crate::download::jpmaqsdownload::download_jpmaqs_indicators_as_df(
+            client_id,
+            client_secret,
+            tickers,
+            metrics,
+            start_date,
+            end_date,
+        )
+        .unwrap(),
+    ))
+}