use alloc::{string::String, vec::Vec};
use std::{
io,
path::{Path, PathBuf},
};
use thiserror::Error;
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum DatasetCompression {
None,
Gzip,
TarGzip,
Zip,
}
#[derive(Debug, Copy, Clone, PartialEq, Eq, Default)]
pub enum CacheMode {
#[default]
UseCache,
Redownload,
}
#[derive(Debug, Copy, Clone, PartialEq, Eq, Default)]
pub enum ArchiveMode {
#[default]
KeepCompressed,
Decompress,
KeepBoth,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DatasetFetchOptions {
pub cache_dir: Option<PathBuf>,
pub cache_mode: CacheMode,
pub archive_mode: ArchiveMode,
}
impl Default for DatasetFetchOptions {
fn default() -> Self {
Self {
cache_dir: None,
cache_mode: CacheMode::UseCache,
archive_mode: ArchiveMode::KeepCompressed,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DatasetArtifact {
pub(crate) dataset_id: &'static str,
pub(crate) path: PathBuf,
pub(crate) compressed_path: Option<PathBuf>,
pub(crate) decompressed_path: Option<PathBuf>,
pub(crate) was_downloaded: bool,
pub(crate) was_decompressed: bool,
}
impl DatasetArtifact {
#[must_use]
pub fn dataset_id(&self) -> &'static str {
self.dataset_id
}
#[must_use]
pub fn path(&self) -> &Path {
&self.path
}
#[must_use]
pub fn compressed_path(&self) -> Option<&Path> {
self.compressed_path.as_deref()
}
#[must_use]
pub fn decompressed_path(&self) -> Option<&Path> {
self.decompressed_path.as_deref()
}
#[must_use]
pub fn was_downloaded(&self) -> bool {
self.was_downloaded
}
#[must_use]
pub fn was_decompressed(&self) -> bool {
self.was_decompressed
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DatasetCollectionArtifact {
pub(crate) dataset_id: &'static str,
pub(crate) paths: Vec<PathBuf>,
pub(crate) compressed_paths: Vec<PathBuf>,
pub(crate) was_downloaded: bool,
pub(crate) was_extracted: bool,
}
impl DatasetCollectionArtifact {
#[must_use]
pub fn dataset_id(&self) -> &'static str {
self.dataset_id
}
#[must_use]
pub fn paths(&self) -> &[PathBuf] {
&self.paths
}
#[must_use]
pub fn compressed_paths(&self) -> &[PathBuf] {
&self.compressed_paths
}
#[must_use]
pub fn was_downloaded(&self) -> bool {
self.was_downloaded
}
#[must_use]
pub fn was_extracted(&self) -> bool {
self.was_extracted
}
}
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub struct DatasetFile {
url: &'static str,
file_name: &'static str,
extracted_file_name: &'static str,
compression: DatasetCompression,
}
impl DatasetFile {
#[must_use]
pub const fn new(
url: &'static str,
file_name: &'static str,
extracted_file_name: &'static str,
compression: DatasetCompression,
) -> Self {
Self { url, file_name, extracted_file_name, compression }
}
#[must_use]
pub fn url(&self) -> &'static str {
self.url
}
#[must_use]
pub fn file_name(&self) -> &'static str {
self.file_name
}
#[must_use]
pub fn extracted_file_name(&self) -> &'static str {
self.extracted_file_name
}
#[must_use]
pub fn compression(&self) -> DatasetCompression {
self.compression
}
}
#[derive(Debug, Error)]
pub enum DatasetError {
#[error("failed to download dataset from {url}: {source}")]
Download {
url: &'static str,
#[source]
source: reqwest::Error,
},
#[error("failed to access dataset path {path}: {source}")]
Io {
path: PathBuf,
#[source]
source: io::Error,
},
#[error("failed to parse dataset {dataset_id} at line {line_number}: {message}")]
Format {
dataset_id: &'static str,
line_number: usize,
message: String,
},
#[error("invalid dataset selection for {dataset_id}: {message}")]
InvalidSelection {
dataset_id: &'static str,
message: String,
},
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct DatasetSmilesRecord {
id: String,
smiles: String,
}
impl DatasetSmilesRecord {
#[must_use]
pub fn new(id: String, smiles: String) -> Self {
Self { id, smiles }
}
#[must_use]
pub fn id(&self) -> &str {
&self.id
}
#[must_use]
pub fn smiles(&self) -> &str {
&self.smiles
}
#[must_use]
pub fn into_smiles(self) -> String {
self.smiles
}
}