Skip to main content

smiles_rs/datasets/
types.rs

1use alloc::{string::String, vec::Vec};
2use std::{
3    io,
4    path::{Path, PathBuf},
5};
6
7use thiserror::Error;
8
9/// Compression used by the upstream dataset artifact.
10#[derive(Debug, Copy, Clone, PartialEq, Eq)]
11pub enum DatasetCompression {
12    /// The upstream file is not compressed.
13    None,
14    /// The upstream file is gzip-compressed.
15    Gzip,
16    /// The upstream file is a gzip-compressed tar archive.
17    TarGzip,
18    /// The upstream file is a zip-compressed archive
19    Zip,
20}
21
22/// Cache behavior for dataset fetches.
23#[derive(Debug, Copy, Clone, PartialEq, Eq, Default)]
24pub enum CacheMode {
25    /// Reuse an existing cached artifact when present.
26    #[default]
27    UseCache,
28    /// Always redownload the upstream artifact.
29    Redownload,
30}
31
32/// How compressed datasets should be materialized locally.
33#[derive(Debug, Copy, Clone, PartialEq, Eq, Default)]
34pub enum ArchiveMode {
35    /// Keep the cached archive file as-is.
36    #[default]
37    KeepCompressed,
38    /// Materialize and return a decompressed copy.
39    Decompress,
40    /// Keep the archive cache and also materialize a decompressed copy.
41    KeepBoth,
42}
43
44/// Options controlling dataset fetch and cache behavior.
45#[derive(Debug, Clone, PartialEq, Eq)]
46pub struct DatasetFetchOptions {
47    /// Override the cache directory. When `None`, a per-user cache directory is
48    /// selected automatically.
49    pub cache_dir: Option<PathBuf>,
50    /// Whether to reuse existing cached files or redownload them.
51    pub cache_mode: CacheMode,
52    /// How compressed datasets should be stored locally.
53    pub archive_mode: ArchiveMode,
54}
55
56impl Default for DatasetFetchOptions {
57    fn default() -> Self {
58        Self {
59            cache_dir: None,
60            cache_mode: CacheMode::UseCache,
61            archive_mode: ArchiveMode::KeepCompressed,
62        }
63    }
64}
65
66/// A materialized dataset artifact on disk.
67#[derive(Debug, Clone, PartialEq, Eq)]
68pub struct DatasetArtifact {
69    pub(crate) dataset_id: &'static str,
70    pub(crate) path: PathBuf,
71    pub(crate) compressed_path: Option<PathBuf>,
72    pub(crate) decompressed_path: Option<PathBuf>,
73    pub(crate) was_downloaded: bool,
74    pub(crate) was_decompressed: bool,
75}
76
77impl DatasetArtifact {
78    /// Returns the dataset identifier that produced this artifact.
79    ///
80    /// # Examples
81    ///
82    /// ```no_run
83    /// use smiles_rs::datasets::{DatasetSource, MASS_SPEC_GYM_SMILES};
84    ///
85    /// let artifact = MASS_SPEC_GYM_SMILES.fetch()?;
86    /// assert_eq!(artifact.dataset_id(), "massspecgym-smiles");
87    /// # Ok::<(), smiles_rs::DatasetError>(())
88    /// ```
89    #[must_use]
90    pub fn dataset_id(&self) -> &'static str {
91        self.dataset_id
92    }
93
94    /// Returns the primary path callers should consume.
95    ///
96    /// This is the decompressed path when one was requested, otherwise the
97    /// cached upstream file path.
98    ///
99    /// # Examples
100    ///
101    /// ```no_run
102    /// use smiles_rs::datasets::{DatasetSource, MASS_SPEC_GYM_SMILES};
103    ///
104    /// let artifact = MASS_SPEC_GYM_SMILES.fetch()?;
105    /// println!("{}", artifact.path().display());
106    /// # Ok::<(), smiles_rs::DatasetError>(())
107    /// ```
108    #[must_use]
109    pub fn path(&self) -> &Path {
110        &self.path
111    }
112
113    /// Returns the cached compressed path, when present.
114    ///
115    /// # Examples
116    ///
117    /// ```no_run
118    /// use smiles_rs::datasets::{ArchiveMode, DatasetFetchOptions, DatasetSource, PUBCHEM_SMILES};
119    ///
120    /// let artifact = PUBCHEM_SMILES.fetch_with_options(&DatasetFetchOptions {
121    ///     archive_mode: ArchiveMode::KeepBoth,
122    ///     ..DatasetFetchOptions::default()
123    /// })?;
124    /// assert!(artifact.compressed_path().is_some());
125    /// # Ok::<(), smiles_rs::DatasetError>(())
126    /// ```
127    #[must_use]
128    pub fn compressed_path(&self) -> Option<&Path> {
129        self.compressed_path.as_deref()
130    }
131
132    /// Returns the cached decompressed path, when present.
133    ///
134    /// # Examples
135    ///
136    /// ```no_run
137    /// use smiles_rs::datasets::{ArchiveMode, DatasetFetchOptions, DatasetSource, PUBCHEM_SMILES};
138    ///
139    /// let artifact = PUBCHEM_SMILES.fetch_with_options(&DatasetFetchOptions {
140    ///     archive_mode: ArchiveMode::Decompress,
141    ///     ..DatasetFetchOptions::default()
142    /// })?;
143    /// assert!(artifact.decompressed_path().is_some());
144    /// # Ok::<(), smiles_rs::DatasetError>(())
145    /// ```
146    #[must_use]
147    pub fn decompressed_path(&self) -> Option<&Path> {
148        self.decompressed_path.as_deref()
149    }
150
151    /// Returns whether the upstream artifact was downloaded during this call.
152    ///
153    /// # Examples
154    ///
155    /// ```no_run
156    /// use smiles_rs::datasets::{
157    ///     CacheMode, DatasetFetchOptions, DatasetSource, MASS_SPEC_GYM_SMILES,
158    /// };
159    ///
160    /// let artifact = MASS_SPEC_GYM_SMILES.fetch_with_options(&DatasetFetchOptions {
161    ///     cache_mode: CacheMode::UseCache,
162    ///     ..DatasetFetchOptions::default()
163    /// })?;
164    /// let _downloaded = artifact.was_downloaded();
165    /// # Ok::<(), smiles_rs::DatasetError>(())
166    /// ```
167    #[must_use]
168    pub fn was_downloaded(&self) -> bool {
169        self.was_downloaded
170    }
171
172    /// Returns whether a decompressed copy was created during this call.
173    ///
174    /// # Examples
175    ///
176    /// ```no_run
177    /// use smiles_rs::datasets::{ArchiveMode, DatasetFetchOptions, DatasetSource, PUBCHEM_SMILES};
178    ///
179    /// let artifact = PUBCHEM_SMILES.fetch_with_options(&DatasetFetchOptions {
180    ///     archive_mode: ArchiveMode::Decompress,
181    ///     ..DatasetFetchOptions::default()
182    /// })?;
183    /// let _decompressed = artifact.was_decompressed();
184    /// # Ok::<(), smiles_rs::DatasetError>(())
185    /// ```
186    #[must_use]
187    pub fn was_decompressed(&self) -> bool {
188        self.was_decompressed
189    }
190}
191
192/// A materialized multi-file dataset collection on disk.
193#[derive(Debug, Clone, PartialEq, Eq)]
194pub struct DatasetCollectionArtifact {
195    pub(crate) dataset_id: &'static str,
196    pub(crate) paths: Vec<PathBuf>,
197    pub(crate) compressed_paths: Vec<PathBuf>,
198    pub(crate) was_downloaded: bool,
199    pub(crate) was_extracted: bool,
200}
201
202impl DatasetCollectionArtifact {
203    /// Returns the dataset identifier that produced this artifact collection.
204    #[must_use]
205    pub fn dataset_id(&self) -> &'static str {
206        self.dataset_id
207    }
208
209    /// Returns the primary paths callers should consume.
210    ///
211    /// For archive-based datasets these are extracted directories when
212    /// extraction was requested, otherwise the cached archive paths.
213    #[must_use]
214    pub fn paths(&self) -> &[PathBuf] {
215        &self.paths
216    }
217
218    /// Returns the cached compressed archive paths.
219    #[must_use]
220    pub fn compressed_paths(&self) -> &[PathBuf] {
221        &self.compressed_paths
222    }
223
224    /// Returns whether any upstream artifact was downloaded during this call.
225    #[must_use]
226    pub fn was_downloaded(&self) -> bool {
227        self.was_downloaded
228    }
229
230    /// Returns whether any upstream archive was extracted during this call.
231    #[must_use]
232    pub fn was_extracted(&self) -> bool {
233        self.was_extracted
234    }
235}
236
237/// Metadata for one file within a downloadable dataset collection.
238#[derive(Debug, Copy, Clone, PartialEq, Eq)]
239pub struct DatasetFile {
240    url: &'static str,
241    file_name: &'static str,
242    extracted_file_name: &'static str,
243    compression: DatasetCompression,
244}
245
246impl DatasetFile {
247    /// Creates static metadata for one downloadable dataset file.
248    #[must_use]
249    pub const fn new(
250        url: &'static str,
251        file_name: &'static str,
252        extracted_file_name: &'static str,
253        compression: DatasetCompression,
254    ) -> Self {
255        Self { url, file_name, extracted_file_name, compression }
256    }
257
258    /// Returns the upstream URL for this dataset file.
259    #[must_use]
260    pub fn url(&self) -> &'static str {
261        self.url
262    }
263
264    /// Returns the cached upstream file name.
265    #[must_use]
266    pub fn file_name(&self) -> &'static str {
267        self.file_name
268    }
269
270    /// Returns the extracted file or directory name.
271    #[must_use]
272    pub fn extracted_file_name(&self) -> &'static str {
273        self.extracted_file_name
274    }
275
276    /// Returns the compression used by this dataset file.
277    #[must_use]
278    pub fn compression(&self) -> DatasetCompression {
279        self.compression
280    }
281}
282
283/// Errors raised while fetching and materializing datasets.
284#[derive(Debug, Error)]
285pub enum DatasetError {
286    /// The HTTP download failed.
287    #[error("failed to download dataset from {url}: {source}")]
288    Download {
289        /// The upstream URL.
290        url: &'static str,
291        /// The underlying transport or HTTP error.
292        #[source]
293        source: reqwest::Error,
294    },
295    /// The on-disk materialization step failed.
296    #[error("failed to access dataset path {path}: {source}")]
297    Io {
298        /// The path being read or written.
299        path: PathBuf,
300        /// The underlying filesystem error.
301        #[source]
302        source: io::Error,
303    },
304    /// The dataset contents did not match the expected record layout.
305    #[error("failed to parse dataset {dataset_id} at line {line_number}: {message}")]
306    Format {
307        /// The dataset identifier.
308        dataset_id: &'static str,
309        /// The 1-based line number within the materialized dataset file.
310        line_number: usize,
311        /// A human-readable explanation of the malformed record.
312        message: String,
313    },
314    /// The requested dataset subset is not valid.
315    #[error("invalid dataset selection for {dataset_id}: {message}")]
316    InvalidSelection {
317        /// The dataset identifier.
318        dataset_id: &'static str,
319        /// A human-readable explanation of the invalid selection.
320        message: String,
321    },
322}
323
324/// One SMILES record from a dataset.
325#[derive(Debug, Clone, PartialEq, Eq)]
326pub struct DatasetSmilesRecord {
327    id: String,
328    smiles: String,
329}
330
331impl DatasetSmilesRecord {
332    /// Creates a dataset SMILES record.
333    #[must_use]
334    pub fn new(id: String, smiles: String) -> Self {
335        Self { id, smiles }
336    }
337
338    /// Returns the dataset-specific record identifier.
339    #[must_use]
340    pub fn id(&self) -> &str {
341        &self.id
342    }
343
344    /// Returns the SMILES string.
345    #[must_use]
346    pub fn smiles(&self) -> &str {
347        &self.smiles
348    }
349
350    /// Consumes the record and returns its SMILES string.
351    #[must_use]
352    pub fn into_smiles(self) -> String {
353        self.smiles
354    }
355}