add Python wrapper for downloading JPMaQS indicators as DataFrame

This commit is contained in:
Palash Tyagi 2024-11-18 01:02:51 +00:00
parent 5c8263d03e
commit 213ea72b20
2 changed files with 81 additions and 0 deletions

View File

@ -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)
}

43
src/py/download.rs Normal file
View File

@ -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(),
))
}