use pyo3::prelude::*;
use std::path::{Path, PathBuf};
use crate::download;
use crate::errors::DownloadError;
fn run_async<F: std::future::Future<Output = Result<(), DownloadError>>>(
fut: F,
) -> PyResult<()> {
let rt = tokio::runtime::Runtime::new()
.map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?;
rt.block_on(fut)
.map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))
}
#[pyclass(skip_from_py_object)]
#[derive(Clone)]
pub struct Client {
threads: usize,
retries: usize,
progress: bool,
}
#[pymethods]
impl Client {
#[new]
#[pyo3(signature = (threads=10, retries=1000, progress=false))]
fn new(threads: usize, retries: usize, progress: bool) -> Self {
Client {
threads,
retries,
progress,
}
}
fn paths(&self, snapshot: &str, data_type: &str) -> PathsDownload {
PathsDownload {
client: self.clone(),
snapshot: snapshot.to_string(),
data_type: data_type.to_string(),
}
}
fn download(&self, path_file: &str) -> DataDownload {
DataDownload {
client: self.clone(),
path_file: PathBuf::from(path_file),
files_only: false,
numbered: false,
strict: false,
}
}
}
#[pyclass]
pub struct PathsDownload {
client: Client,
snapshot: String,
data_type: String,
}
#[pymethods]
impl PathsDownload {
fn to(&self, dst: &str) -> PyResult<()> {
let dst_path = Path::new(dst);
std::fs::create_dir_all(dst_path)
.map_err(|e| pyo3::exceptions::PyIOError::new_err(e.to_string()))?;
let options = download::DownloadOptions {
snapshot: self.snapshot.clone(),
data_type: &self.data_type,
dst: dst_path,
max_retries: self.client.retries,
progress: self.client.progress,
..Default::default()
};
run_async(download::download_paths(options))
}
}
#[pyclass]
pub struct DataDownload {
client: Client,
path_file: PathBuf,
files_only: bool,
numbered: bool,
strict: bool,
}
#[pymethods]
impl DataDownload {
fn files_only(mut slf: PyRefMut<'_, Self>) -> PyRefMut<'_, Self> {
slf.files_only = true;
slf
}
fn numbered(mut slf: PyRefMut<'_, Self>) -> PyRefMut<'_, Self> {
slf.numbered = true;
slf
}
fn strict(mut slf: PyRefMut<'_, Self>) -> PyRefMut<'_, Self> {
slf.strict = true;
slf
}
fn to(&self, dst: &str) -> PyResult<()> {
let dst_path = Path::new(dst);
std::fs::create_dir_all(dst_path)
.map_err(|e| pyo3::exceptions::PyIOError::new_err(e.to_string()))?;
let options = download::DownloadOptions {
paths: &self.path_file,
dst: dst_path,
threads: self.client.threads,
max_retries: self.client.retries,
numbered: self.numbered,
files_only: self.files_only,
progress: self.client.progress,
strict: self.strict,
..Default::default()
};
run_async(download::download(options))
}
}
#[pymodule]
pub fn ccdown(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<Client>()?;
m.add_class::<PathsDownload>()?;
m.add_class::<DataDownload>()?;
Ok(())
}