dataset_core/utils.rs
1use crate::DatasetError;
2use flate2::read::GzDecoder;
3use sha2::{Digest, Sha256};
4use std::fmt::Write;
5use std::fs::File;
6use std::io;
7use std::io::Read;
8use std::path::{Path, PathBuf};
9use tar::Archive;
10use zip::ZipArchive;
11use zip::result::ZipError;
12
13/// Download a remote file into the given directory.
14///
15/// It downloads the content at `url` into `storage_path`, using the [`ureq`] crate. It names the
16/// file after the last segment of the URL, unless the caller supplies a custom filename.
17///
18/// When the filename comes from the URL, it strips any trailing `?query` or `#fragment` (for
19/// example, `.../iris.csv?raw=1` yields `iris.csv`). A URL that ends in `/` has no final segment,
20/// so the caller must supply a custom `filename` in that case.
21///
22/// # Parameters
23///
24/// - `url` - The URL to download.
25/// - `storage_path` - The directory to store the downloaded file in.
26/// - `filename` - Optional custom filename (with extension). If `None`, the filename comes from
27/// the last segment of the URL.
28///
29/// # Errors
30///
31/// - `DatasetError` - Returned when the download fails, or when the function cannot derive the
32/// filename from the URL.
33///
34/// # Example
35/// ```no_run
36/// use dataset_core::download_to;
37/// use std::path::Path;
38///
39/// let download_dir = "./download_example";
40/// std::fs::create_dir_all(download_dir).unwrap();
41///
42/// // Download a file from the internet
43/// let url = "https://gist.githubusercontent.com/curran/a08a1080b88344b0c8a7/raw/0e7a9b0a5d22642a06d3d5b9bcbad9890c8ee534/iris.csv";
44///
45/// // Use filename from URL
46/// download_to(url, Path::new(download_dir), None).unwrap();
47/// assert!(Path::new(download_dir).join("iris.csv").exists());
48///
49/// // Use custom filename
50/// download_to(url, Path::new(download_dir), Some("custom.csv")).unwrap();
51/// assert!(Path::new(download_dir).join("custom.csv").exists());
52/// ```
53pub fn download_to(
54 url: &str,
55 storage_path: &Path,
56 filename: Option<&str>,
57) -> Result<(), DatasetError> {
58 let filename = match filename {
59 Some(name) => name,
60 None => filename_from_url(url).ok_or_else(|| {
61 DatasetError::ValidationError(
62 "Invalid URL: cannot extract filename from URL".to_string(),
63 )
64 })?,
65 };
66
67 let save_path = storage_path.join(filename);
68
69 let mut response = ureq::get(url).call()?;
70 let mut body = response.body_mut().as_reader();
71
72 let mut file = File::create(save_path)?;
73 io::copy(&mut body, &mut file)?;
74
75 Ok(())
76}
77
78/// Download a remote file into the given directory. It retries transient failures.
79///
80/// This function wraps [`download_to`] for the unreliable hosts where many public datasets live.
81/// If the download fails, it retries up to `retries` more times. Each attempt waits twice as
82/// long as the last (500 ms, then 1 s, 2 s, and so on). When the caller sets `retries` to `0`,
83/// this function behaves exactly like [`download_to`].
84///
85/// A retry cannot fix two kinds of failure: a filename the function cannot derive from the URL,
86/// and a local file the function cannot create. The function returns either failure right away,
87/// without a retry. Once every attempt fails, it returns the last download error.
88///
89/// # Parameters
90///
91/// - `url` - The URL to download.
92/// - `storage_path` - The directory to store the downloaded file in.
93/// - `filename` - Optional custom filename (with extension). If `None`, the filename comes from
94/// the last segment of the URL.
95/// - `retries` - How many **additional** attempts to make after the first one fails.
96///
97/// # Errors
98///
99/// - `DatasetError` - Returned when every attempt fails (it returns the last error), or right
100/// away for an error a retry cannot fix.
101///
102/// # Example
103/// ```no_run
104/// use dataset_core::download_to_with_retries;
105/// use std::path::Path;
106///
107/// let download_dir = Path::new("./download_retry_example");
108/// std::fs::create_dir_all(download_dir).unwrap();
109///
110/// // Try up to three times in total before giving up.
111/// download_to_with_retries(
112/// "https://archive.ics.uci.edu/static/public/53/iris.zip",
113/// download_dir,
114/// Some("iris.zip"),
115/// 2,
116/// )
117/// .unwrap();
118/// ```
119pub fn download_to_with_retries(
120 url: &str,
121 storage_path: &Path,
122 filename: Option<&str>,
123 retries: u32,
124) -> Result<(), DatasetError> {
125 let mut attempt = 0;
126
127 loop {
128 match download_to(url, storage_path, filename) {
129 Ok(()) => return Ok(()),
130 // A retry cannot fix a malformed URL or an unwritable target.
131 Err(e @ (DatasetError::ValidationError(_) | DatasetError::IoError(_))) => {
132 return Err(e);
133 }
134 Err(e) => {
135 if attempt >= retries {
136 return Err(e);
137 }
138 std::thread::sleep(RETRY_BASE_DELAY * 2u32.pow(attempt));
139 attempt += 1;
140 }
141 }
142 }
143}
144
145/// How long [`download_to_with_retries`] waits before its first retry. Each further retry
146/// doubles this wait.
147const RETRY_BASE_DELAY: std::time::Duration = std::time::Duration::from_millis(500);
148
149/// Derive a filename from the last path segment of a URL.
150///
151/// It strips any `?query` / `#fragment` suffix, and returns `None` when the resulting segment
152/// is empty (for example, when the URL ends in `/`).
153fn filename_from_url(url: &str) -> Option<&str> {
154 // `rsplit('/').next()` yields the whole string when there is no '/'.
155 let last_segment = url.rsplit('/').next()?;
156 let name = last_segment
157 .split(['?', '#'])
158 .next()
159 .unwrap_or(last_segment);
160 (!name.is_empty()).then_some(name)
161}
162
163/// Extract a zip archive into a target directory.
164///
165/// It uses [`ZipArchive`] from the [`zip`] crate.
166///
167/// # Parameters
168///
169/// - `file_path` - Path to the `.zip` file to extract.
170/// - `extract_dir` - Directory to extract the archive contents into.
171///
172/// # Errors
173///
174/// - `DatasetError` - Returned when opening the zip file fails or when extraction fails.
175///
176/// # Example
177/// ```no_run
178/// use dataset_core::unzip;
179/// use std::fs::File;
180/// use std::io::Write;
181/// use std::path::Path;
182/// use zip::write::SimpleFileOptions;
183///
184/// let work_dir = Path::new("./unzip_example");
185/// std::fs::create_dir_all(work_dir).unwrap();
186///
187/// // Build a small `.zip` containing a single `hello.txt` entry.
188/// let archive_path = work_dir.join("data.zip");
189/// let mut zip = zip::ZipWriter::new(File::create(&archive_path).unwrap());
190/// zip.start_file("hello.txt", SimpleFileOptions::default()).unwrap();
191/// zip.write_all(b"hello world").unwrap();
192/// zip.finish().unwrap();
193///
194/// // Extract it into `work_dir`.
195/// unzip(&archive_path, work_dir).unwrap();
196/// assert!(work_dir.join("hello.txt").exists());
197/// ```
198pub fn unzip(file_path: &Path, extract_dir: &Path) -> Result<(), DatasetError> {
199 let file = File::open(file_path).map_err(|e| DatasetError::from(ZipError::Io(e)))?;
200
201 ZipArchive::new(file)?.extract(extract_dir)?;
202
203 Ok(())
204}
205
206/// Decompress a gzip (`.gz`) file into a single output file.
207///
208/// Unlike [`unzip`], which extracts a multi-entry archive into a directory, a gzip stream wraps
209/// exactly **one** file. As a result, this function writes the decompressed bytes to a single
210/// `output_path`. It streams the bytes through [`flate2::read::GzDecoder`], so it never holds
211/// the whole file in memory at once. This suits large datasets such as the gzip-compressed
212/// `covtype.data.gz`.
213///
214/// This function creates the output file, or truncates it if the file already exists. Any
215/// leading directories in `output_path` must already exist.
216///
217/// # Parameters
218///
219/// - `file_path` - Path to the `.gz` file to decompress.
220/// - `output_path` - Path of the decompressed file to write (including filename).
221///
222/// # Errors
223///
224/// - `DatasetError::IoError` - Returned when opening the source fails, the gzip
225/// stream is malformed, or writing the output fails.
226///
227/// # Example
228/// ```no_run
229/// use dataset_core::{download_to, gunzip};
230/// use std::path::Path;
231///
232/// let work_dir = Path::new("./gunzip_example");
233/// std::fs::create_dir_all(work_dir).unwrap();
234///
235/// // Download a gzip-compressed dataset. Then decompress it in place.
236/// download_to("https://example.com/data.csv.gz", work_dir, Some("data.csv.gz")).unwrap();
237/// gunzip(&work_dir.join("data.csv.gz"), &work_dir.join("data.csv")).unwrap();
238/// assert!(work_dir.join("data.csv").exists());
239/// ```
240pub fn gunzip(file_path: &Path, output_path: &Path) -> Result<(), DatasetError> {
241 let input = File::open(file_path)?;
242 let mut decoder = GzDecoder::new(input);
243 let mut output = File::create(output_path)?;
244 io::copy(&mut decoder, &mut output)?;
245
246 Ok(())
247}
248
249/// Extract a tar (`.tar`) archive into a target directory.
250///
251/// It uses [`tar::Archive`] for the extraction.
252///
253/// This is the tar analogue of [`unzip`]: a `.tar` bundles a whole directory tree (unlike
254/// [`gunzip`], which decompresses a single-file gzip stream). For a common gzip-compressed
255/// tarball (`.tar.gz` / `.tgz`), use [`untar_gz`] instead. It streams the decompression and
256/// extraction together, and never writes an intermediate `.tar` file to disk.
257///
258/// This function unpacks the archive's entries **relative to** `extract_dir`, and creates that
259/// directory if needed. The [`tar`] crate rejects any entry whose path would escape
260/// `extract_dir`.
261///
262/// # Parameters
263///
264/// - `file_path` - Path to the `.tar` file to extract.
265/// - `extract_dir` - Directory to extract the archive contents into.
266///
267/// # Errors
268///
269/// - `DatasetError::IoError` - Returned when opening the archive fails or when
270/// extraction fails (a malformed archive, or an entry that cannot be written).
271///
272/// # Example
273/// ```no_run
274/// use dataset_core::{download_to, untar};
275/// use std::path::Path;
276///
277/// let work_dir = Path::new("./untar_example");
278/// std::fs::create_dir_all(work_dir).unwrap();
279///
280/// // Download a tar archive. Then extract it in place.
281/// download_to("https://example.com/data.tar", work_dir, Some("data.tar")).unwrap();
282/// untar(&work_dir.join("data.tar"), work_dir).unwrap();
283/// ```
284pub fn untar(file_path: &Path, extract_dir: &Path) -> Result<(), DatasetError> {
285 let file = File::open(file_path)?;
286 Archive::new(file).unpack(extract_dir)?;
287
288 Ok(())
289}
290
291/// Extract a gzip-compressed tar (`.tar.gz` / `.tgz`) archive into a target directory.
292///
293/// This function combines the two layers of a gzipped tarball in one streaming pass. The bytes
294/// flow through [`flate2::read::GzDecoder`] (the gzip layer, as in [`gunzip`]), straight into
295/// [`tar::Archive`] (the tar layer, as in [`untar`]). This function never writes the intermediate
296/// uncompressed `.tar` to disk, so it suits large datasets distributed as `.tar.gz`.
297///
298/// This function unpacks the archive's entries **relative to** `extract_dir`, and creates that
299/// directory if needed. The [`tar`] crate rejects any entry whose path would escape
300/// `extract_dir`.
301///
302/// # Parameters
303///
304/// - `file_path` - Path to the `.tar.gz` (or `.tgz`) file to extract.
305/// - `extract_dir` - Directory to extract the archive contents into.
306///
307/// # Errors
308///
309/// - `DatasetError::IoError` - Returned when opening the source fails, the gzip
310/// stream is malformed, or the tar extraction fails.
311///
312/// # Example
313/// ```no_run
314/// use dataset_core::{download_to, untar_gz};
315/// use std::path::Path;
316///
317/// let work_dir = Path::new("./untar_gz_example");
318/// std::fs::create_dir_all(work_dir).unwrap();
319///
320/// // Download a gzip-compressed tarball. Then extract it in place.
321/// download_to("https://example.com/data.tar.gz", work_dir, Some("data.tar.gz")).unwrap();
322/// untar_gz(&work_dir.join("data.tar.gz"), work_dir).unwrap();
323/// ```
324pub fn untar_gz(file_path: &Path, extract_dir: &Path) -> Result<(), DatasetError> {
325 let input = File::open(file_path)?;
326 let decoder = GzDecoder::new(input);
327 Archive::new(decoder).unpack(extract_dir)?;
328
329 Ok(())
330}
331
332/// Create a temporary directory under the given parent directory.
333///
334/// This is a small wrapper around [`tempfile::Builder`]. [`acquire_dataset`] uses it internally
335/// to keep intermediate download and extraction artifacts isolated. The [`tempfile::TempDir`]
336/// removes the directory automatically when it is dropped.
337fn create_temp_dir(tempdir_in: &Path) -> Result<tempfile::TempDir, DatasetError> {
338 let temp_dir = tempfile::Builder::new().tempdir_in(tempdir_in)?;
339
340 Ok(temp_dir)
341}
342
343/// Compute a file's SHA256 hash and return it as a lowercase hex string.
344///
345/// This function streams the file in 8 KiB chunks, so hashing a multi-gigabyte dataset costs
346/// no more memory than hashing a small one.
347///
348/// This is the helper to reach for when **pinning** a hash. Run it once against a freshly
349/// downloaded file. Paste the result into the `expected_sha256` value that you pass to
350/// [`acquire_dataset`]. To check a file against a hash you already have, use [`verify_sha256`]
351/// instead of comparing strings yourself.
352///
353/// # Parameters
354///
355/// - `path` - Path to the file to hash.
356///
357/// # Returns
358///
359/// - `String` - The SHA256 digest as 64 lowercase hex characters.
360///
361/// # Errors
362///
363/// - `DatasetError::IoError` - Returned when the file cannot be opened or read.
364///
365/// # Example
366/// ```no_run
367/// use dataset_core::sha256_file;
368/// use std::path::Path;
369///
370/// let digest = sha256_file(Path::new("./data/iris.csv")).unwrap();
371/// println!("pin this: {digest}");
372/// ```
373pub fn sha256_file(path: &Path) -> Result<String, DatasetError> {
374 let mut file = File::open(path)?;
375
376 let mut hasher = Sha256::new();
377 let mut buf = [0u8; 8192];
378
379 loop {
380 let read = file.read(&mut buf)?;
381 if read == 0 {
382 break;
383 }
384 hasher.update(&buf[..read]);
385 }
386
387 let digest = hasher.finalize();
388 let mut hex = String::with_capacity(digest.len() * 2);
389 for b in digest {
390 // A `write!` into a `String` cannot fail.
391 let _ = write!(hex, "{:02x}", b);
392 }
393
394 Ok(hex)
395}
396
397/// Verify that a file's SHA256 hash matches an expected value (case-insensitive).
398///
399/// This is the same check that [`acquire_dataset`] performs internally on cached and freshly
400/// prepared files. Callers that need to validate a file outside that workflow can use it too.
401/// Most commonly, this is a test that asserts the file on disk is the expected one. This
402/// function returns `true` when the computed hash matches `expected_hex`.
403///
404/// # Parameters
405///
406/// - `path` - Path to the file to verify.
407/// - `expected_hex` - The expected SHA256 digest, in hex (either case).
408///
409/// # Returns
410///
411/// - `bool` - `true` if the file's hash matches `expected_hex`, `false` otherwise.
412///
413/// # Errors
414///
415/// - `DatasetError::IoError` - Returned when the file cannot be opened or read.
416///
417/// # Example
418/// ```no_run
419/// use dataset_core::verify_sha256;
420/// use std::path::Path;
421///
422/// const IRIS_SHA256: &str = "c52742e50315a99f956a383faedf7575552675f6409ef0f9a47076dd08479930";
423///
424/// assert!(verify_sha256(Path::new("./data/iris.csv"), IRIS_SHA256).unwrap());
425/// ```
426pub fn verify_sha256(path: &Path, expected_hex: &str) -> Result<bool, DatasetError> {
427 Ok(sha256_file(path)?.eq_ignore_ascii_case(expected_hex))
428}
429
430/// Read a file as Latin-1 (ISO-8859-1) text.
431///
432/// This function maps every byte to the Unicode scalar with the same value. That is exactly
433/// what Latin-1 decoding means, and what scikit-learn does for the older text corpora. Unlike
434/// [`std::fs::read_to_string`], this function never fails on non-UTF-8 input, and it never
435/// replaces bytes with `U+FFFD`. The decoding is lossless and reversible, so a corpus whose
436/// encoding is unknown or mixed survives the round trip.
437///
438/// Use it for raw document collections (newsgroup posts, movie reviews, and so on) that predate
439/// UTF-8. For data you know is UTF-8, prefer `std::fs::read_to_string`.
440///
441/// # Parameters
442///
443/// - `path` - Path to the file to read.
444///
445/// # Returns
446///
447/// - `String` - The file's contents, decoded byte-for-byte as Latin-1.
448///
449/// # Errors
450///
451/// - `DatasetError::IoError` - Returned when the file cannot be opened or read.
452///
453/// # Example
454/// ```no_run
455/// use dataset_core::read_latin1;
456/// use std::path::Path;
457///
458/// // A byte that is invalid UTF-8 decodes to the matching Latin-1 character
459/// // instead of failing the read.
460/// let text = read_latin1(Path::new("./corpus/post_00001")).unwrap();
461/// assert!(!text.is_empty());
462/// ```
463pub fn read_latin1(path: &Path) -> Result<String, DatasetError> {
464 let bytes = std::fs::read(path)?;
465
466 Ok(bytes.iter().map(|&b| b as char).collect())
467}
468
469/// State of the destination file for the dataset to cache.
470enum CacheState {
471 /// Destination file exists, and its hash matches (if a hash was given). The code reuses
472 /// it without changes.
473 Fresh,
474 /// Destination exists, but its hash does not match. The code must replace it.
475 Stale,
476 /// Destination does not exist. The code must prepare a new file.
477 Missing,
478}
479
480/// Make sure `dir` exists, then classify the destination file `dst`.
481///
482/// When `expected_sha256` is `None`, any existing file counts as [`CacheState::Fresh`].
483fn inspect_cache(
484 dir: &Path,
485 dst: &Path,
486 expected_sha256: Option<&str>,
487) -> Result<CacheState, DatasetError> {
488 if !dir.exists() {
489 std::fs::create_dir_all(dir)?;
490 }
491
492 if !dst.exists() {
493 return Ok(CacheState::Missing);
494 }
495
496 match expected_sha256 {
497 Some(hash) if !verify_sha256(dst, hash)? => Ok(CacheState::Stale),
498 _ => Ok(CacheState::Fresh),
499 }
500}
501
502/// Get a dataset file, using a preparation closure that the caller provides.
503///
504/// This is the single entry point for getting a dataset, and the recommended way to populate a
505/// storage directory. It checks whether it can reuse the destination file, and creates a
506/// temporary directory when a new file is needed. It then delegates file preparation to a
507/// closure that the caller provides. It optionally validates the prepared file with SHA256, and
508/// atomically moves it to the final destination.
509///
510/// The function itself does not perform network I/O. The `prepare_file` closure prepares the
511/// dataset file. This may include downloading the file, extracting archives, or locating files
512/// inside an extracted directory.
513///
514/// # Parameters
515///
516/// - `dir` - Target storage directory path.
517/// - `filename` - Final dataset filename (stored as `dir/filename`).
518/// Include the file extension (for example, `"iris.csv"`).
519/// - `dataset_name` - Dataset name for error messages (for example, `"iris"`).
520/// - `expected_sha256` - Optional expected SHA256 hash of the dataset file. If `None`,
521/// the function accepts any existing file at the destination without validation, and
522/// newly prepared files skip SHA256 verification.
523/// - `prepare_file` - Closure that prepares the dataset file in the temporary directory.
524/// - Input: `temp_dir: &Path` - Path to the temporary directory. The caller should do
525/// file work inside this directory: the function cleans it up automatically when the
526/// closure returns. This is not a requirement. (The function moves the file to the
527/// final destination. It does not copy the file.)
528/// - Output: `Result<PathBuf, DatasetError>` - Path to the prepared dataset file (the
529/// function moves this file to `dir/filename`).
530/// - Responsibility: This closure can do any work needed to prepare the dataset file.
531/// This can include downloading it (use [`download_to`] from this crate), extracting
532/// archives (use [`unzip`] from this crate), or locating files inside extracted
533/// folders. The returned `PathBuf` must point to the final dataset file, ready for
534/// validation.
535///
536/// # Returns
537///
538/// - `PathBuf` - Path to the final dataset file (`dir/filename`).
539///
540/// # Errors
541///
542/// - `DatasetError::IoError` - Returned when directory creation, file operations, or
543/// hash verification fails.
544/// - `DatasetError::Sha256ValidationFailed` - Returned when the caller provides
545/// `expected_sha256` and the prepared file's SHA256 hash does not match it.
546/// - Any error the `prepare_file` closure returns.
547///
548/// # Example
549/// ```no_run
550/// // Implement the file preparation process for the Iris dataset.
551///
552/// /// The URL for the Iris dataset.
553/// ///
554/// /// # Citation
555/// ///
556/// /// R. A. Fisher. "Iris," UCI Machine Learning Repository, \[Online\].
557/// /// Available: <https://doi.org/10.24432/C56C76>
558/// const IRIS_DATA_URL: &str = "https://gist.githubusercontent.com/curran/a08a1080b88344b0c8a7/raw/0e7a9b0a5d22642a06d3d5b9bcbad9890c8ee534/iris.csv";
559///
560/// /// The name of the Iris dataset file.
561/// const IRIS_FILENAME: &str = "iris.csv";
562///
563/// /// The SHA256 hash of the Iris dataset file.
564/// const IRIS_SHA256: &str = "c52742e50315a99f956a383faedf7575552675f6409ef0f9a47076dd08479930";
565///
566/// /// The name of the dataset.
567/// const IRIS_DATASET_NAME: &str = "iris";
568///
569/// use dataset_core::acquire_dataset;
570/// use dataset_core::download_to;
571///
572/// fn main() {
573/// let dir = "./somewhere";
574///
575/// let file_path = acquire_dataset(
576/// // Target storage directory path
577/// dir,
578/// // Final dataset filename (stored as `dir/filename`)
579/// IRIS_FILENAME,
580/// // Dataset name for error messages
581/// IRIS_DATASET_NAME,
582/// // Expected SHA256 hash of the dataset file
583/// Some(IRIS_SHA256),
584/// // Closure that prepares the dataset file in the temporary directory
585/// |temp_path| {
586/// // Download the dataset into the temporary directory
587/// download_to(IRIS_DATA_URL, temp_path, None)?;
588/// Ok(temp_path.join(IRIS_FILENAME))
589/// },
590/// ).unwrap();
591///
592/// // `file_path` is now the path to the Iris dataset file.
593/// // Use it to locate or parse the dataset.
594/// }
595/// ```
596pub fn acquire_dataset<F>(
597 dir: &str,
598 filename: &str,
599 dataset_name: &str,
600 expected_sha256: Option<&str>,
601 prepare_file: F,
602) -> Result<PathBuf, DatasetError>
603where
604 F: FnOnce(&Path) -> Result<PathBuf, DatasetError>,
605{
606 let dir_path = Path::new(dir);
607 let dst = dir_path.join(filename);
608
609 // Reuse a valid cached file without invoking the preparation closure.
610 let state = inspect_cache(dir_path, &dst, expected_sha256)?;
611 if matches!(state, CacheState::Fresh) {
612 return Ok(dst);
613 }
614
615 // Prepare the new file inside a temp dir. The temp dir cleans up on drop,
616 // including when the code returns early below.
617 let temp_dir = create_temp_dir(dir_path)?;
618 let src = prepare_file(temp_dir.path())?;
619
620 // Validate the freshly prepared file before it lands at the final path.
621 if let Some(hash) = expected_sha256
622 && !verify_sha256(&src, hash)?
623 {
624 return Err(DatasetError::sha256_validation_failed(
625 dataset_name,
626 filename,
627 ));
628 }
629
630 // The code must remove a stale file first: `fs::rename` does not overwrite on all platforms.
631 if matches!(state, CacheState::Stale) {
632 std::fs::remove_file(&dst)?;
633 }
634 std::fs::rename(&src, &dst)?;
635
636 Ok(dst)
637}
638
639#[cfg(test)]
640mod tests {
641 use super::*;
642 use std::fs::{self, File, create_dir_all, remove_dir_all};
643 use std::io::Write;
644
645 /// SHA256 of "hello world"
646 const HELLO_WORLD_SHA256: &str =
647 "b94d27b9934d3e08a52e52d7da7dabfac484efe37a5380ee9088f7ace2efcde9";
648
649 /// SHA256 of an empty file
650 const EMPTY_SHA256: &str = "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855";
651
652 /// All-zero hash (always wrong)
653 const ZERO_SHA256: &str = "0000000000000000000000000000000000000000000000000000000000000000";
654
655 #[test]
656 fn create_temp_dir_returns_existing_path() {
657 let parent = "./test_create_temp_dir_returns_existing_path";
658 create_dir_all(parent).unwrap();
659
660 let temp_dir = create_temp_dir(Path::new(parent)).unwrap();
661 assert!(temp_dir.path().exists());
662
663 remove_dir_all(parent).unwrap();
664 }
665
666 #[test]
667 fn create_temp_dir_cleanup_on_drop() {
668 let parent = "./test_create_temp_dir_cleanup_on_drop";
669 create_dir_all(parent).unwrap();
670
671 let temp_dir = create_temp_dir(Path::new(parent)).unwrap();
672 let temp_path = temp_dir.path().to_path_buf();
673
674 assert!(temp_path.exists());
675 drop(temp_dir);
676 assert!(!temp_path.exists());
677
678 remove_dir_all(parent).unwrap();
679 }
680
681 #[test]
682 fn create_temp_dir_nonexistent_parent_errors() {
683 let result = create_temp_dir(Path::new("./nonexistent_parent_xyz_abc_123"));
684 assert!(result.is_err());
685 }
686
687 /// Compress `content` into a gzip file at `path` (test helper).
688 fn write_gz(path: &Path, content: &[u8]) {
689 use flate2::Compression;
690 use flate2::write::GzEncoder;
691
692 let file = File::create(path).unwrap();
693 let mut encoder = GzEncoder::new(file, Compression::default());
694 encoder.write_all(content).unwrap();
695 encoder.finish().unwrap();
696 }
697
698 #[test]
699 fn gunzip_round_trips_content() {
700 let dir = "./test_gunzip_round_trips_content";
701 create_dir_all(dir).unwrap();
702 let dir_path = Path::new(dir);
703
704 let payload = b"col_a,col_b\n1,2\n3,4\n";
705 let gz_path = dir_path.join("data.csv.gz");
706 let out_path = dir_path.join("data.csv");
707 write_gz(&gz_path, payload);
708
709 gunzip(&gz_path, &out_path).unwrap();
710
711 assert_eq!(fs::read(&out_path).unwrap(), payload);
712
713 remove_dir_all(dir).unwrap();
714 }
715
716 #[test]
717 fn gunzip_nonexistent_source_errors() {
718 let result = gunzip(
719 Path::new("./no_such_file_for_gunzip_test.gz"),
720 Path::new("./no_such_file_for_gunzip_test.out"),
721 );
722 assert!(result.is_err());
723 }
724
725 /// Build a tar archive at `path` from `(name, content)` entries (test helper).
726 fn write_tar(path: &Path, entries: &[(&str, &[u8])]) {
727 let file = File::create(path).unwrap();
728 let mut builder = tar::Builder::new(file);
729 for (name, content) in entries {
730 let mut header = tar::Header::new_gnu();
731 header.set_size(content.len() as u64);
732 header.set_mode(0o644);
733 header.set_cksum();
734 builder.append_data(&mut header, name, *content).unwrap();
735 }
736 builder.into_inner().unwrap().sync_all().unwrap();
737 }
738
739 /// Build a gzip-compressed tar archive at `path` (test helper).
740 fn write_tar_gz(path: &Path, entries: &[(&str, &[u8])]) {
741 use flate2::Compression;
742 use flate2::write::GzEncoder;
743
744 let file = File::create(path).unwrap();
745 let encoder = GzEncoder::new(file, Compression::default());
746 let mut builder = tar::Builder::new(encoder);
747 for (name, content) in entries {
748 let mut header = tar::Header::new_gnu();
749 header.set_size(content.len() as u64);
750 header.set_mode(0o644);
751 header.set_cksum();
752 builder.append_data(&mut header, name, *content).unwrap();
753 }
754 builder.into_inner().unwrap().finish().unwrap();
755 }
756
757 #[test]
758 fn untar_round_trips_entries() {
759 let dir = "./test_untar_round_trips_entries";
760 create_dir_all(dir).unwrap();
761 let dir_path = Path::new(dir);
762
763 let tar_path = dir_path.join("data.tar");
764 write_tar(
765 &tar_path,
766 &[("a.txt", b"hello"), ("nested/b.txt", b"world")],
767 );
768
769 let out_dir = dir_path.join("extracted");
770 untar(&tar_path, &out_dir).unwrap();
771
772 assert_eq!(fs::read(out_dir.join("a.txt")).unwrap(), b"hello");
773 assert_eq!(fs::read(out_dir.join("nested/b.txt")).unwrap(), b"world");
774
775 remove_dir_all(dir).unwrap();
776 }
777
778 #[test]
779 fn untar_nonexistent_source_errors() {
780 let result = untar(
781 Path::new("./no_such_file_for_untar_test.tar"),
782 Path::new("./no_such_dir_for_untar_test"),
783 );
784 assert!(result.is_err());
785 }
786
787 #[test]
788 fn untar_gz_round_trips_entries() {
789 let dir = "./test_untar_gz_round_trips_entries";
790 create_dir_all(dir).unwrap();
791 let dir_path = Path::new(dir);
792
793 let tar_gz_path = dir_path.join("data.tar.gz");
794 write_tar_gz(
795 &tar_gz_path,
796 &[("a.txt", b"hello"), ("nested/b.txt", b"world")],
797 );
798
799 let out_dir = dir_path.join("extracted");
800 untar_gz(&tar_gz_path, &out_dir).unwrap();
801
802 assert_eq!(fs::read(out_dir.join("a.txt")).unwrap(), b"hello");
803 assert_eq!(fs::read(out_dir.join("nested/b.txt")).unwrap(), b"world");
804
805 remove_dir_all(dir).unwrap();
806 }
807
808 #[test]
809 fn untar_gz_nonexistent_source_errors() {
810 let result = untar_gz(
811 Path::new("./no_such_file_for_untar_gz_test.tar.gz"),
812 Path::new("./no_such_dir_for_untar_gz_test"),
813 );
814 assert!(result.is_err());
815 }
816
817 #[test]
818 fn verify_sha256_correct_hash() {
819 let dir = "./test_verify_sha256_correct_hash";
820 create_dir_all(dir).unwrap();
821 let path = Path::new(dir).join("f.txt");
822 File::create(&path)
823 .unwrap()
824 .write_all(b"hello world")
825 .unwrap();
826
827 assert!(verify_sha256(&path, HELLO_WORLD_SHA256).unwrap());
828
829 remove_dir_all(dir).unwrap();
830 }
831
832 #[test]
833 fn verify_sha256_uppercase_hash() {
834 let dir = "./test_verify_sha256_uppercase_hash";
835 create_dir_all(dir).unwrap();
836 let path = Path::new(dir).join("f.txt");
837 File::create(&path)
838 .unwrap()
839 .write_all(b"hello world")
840 .unwrap();
841
842 assert!(verify_sha256(&path, &HELLO_WORLD_SHA256.to_uppercase()).unwrap());
843
844 remove_dir_all(dir).unwrap();
845 }
846
847 #[test]
848 fn verify_sha256_wrong_hash_returns_false() {
849 let dir = "./test_verify_sha256_wrong_hash_returns_false";
850 create_dir_all(dir).unwrap();
851 let path = Path::new(dir).join("f.txt");
852 File::create(&path)
853 .unwrap()
854 .write_all(b"hello world")
855 .unwrap();
856
857 assert!(!verify_sha256(&path, ZERO_SHA256).unwrap());
858
859 remove_dir_all(dir).unwrap();
860 }
861
862 #[test]
863 fn verify_sha256_empty_file() {
864 let dir = "./test_verify_sha256_empty_file";
865 create_dir_all(dir).unwrap();
866 let path = Path::new(dir).join("empty.txt");
867 File::create(&path).unwrap();
868
869 assert!(verify_sha256(&path, EMPTY_SHA256).unwrap());
870
871 remove_dir_all(dir).unwrap();
872 }
873
874 #[test]
875 fn verify_sha256_nonexistent_file_errors() {
876 let result = verify_sha256(Path::new("./no_such_file_sha256_test.txt"), ZERO_SHA256);
877 assert!(result.is_err());
878 }
879
880 #[test]
881 fn sha256_file_returns_lowercase_hex_digest() {
882 let dir = "./test_sha256_file_returns_lowercase_hex_digest";
883 create_dir_all(dir).unwrap();
884 let path = Path::new(dir).join("f.txt");
885 File::create(&path)
886 .unwrap()
887 .write_all(b"hello world")
888 .unwrap();
889
890 let digest = sha256_file(&path).unwrap();
891 assert_eq!(digest, HELLO_WORLD_SHA256);
892 assert_eq!(digest.len(), 64);
893 assert!(
894 digest
895 .chars()
896 .all(|c| c.is_ascii_hexdigit() && !c.is_ascii_uppercase())
897 );
898
899 remove_dir_all(dir).unwrap();
900 }
901
902 #[test]
903 fn sha256_file_hashes_larger_than_one_chunk() {
904 let dir = "./test_sha256_file_hashes_larger_than_one_chunk";
905 create_dir_all(dir).unwrap();
906 let path = Path::new(dir).join("big.bin");
907 // Larger than the 8 KiB read buffer, so the streaming loop runs several times.
908 let payload = vec![0xABu8; 8192 * 3 + 17];
909 fs::write(&path, &payload).unwrap();
910
911 // The streamed digest must equal the one-shot digest of the same bytes.
912 let expected = {
913 let mut hasher = Sha256::new();
914 hasher.update(&payload);
915 hasher
916 .finalize()
917 .iter()
918 .map(|b| format!("{:02x}", b))
919 .collect::<String>()
920 };
921 assert_eq!(sha256_file(&path).unwrap(), expected);
922
923 remove_dir_all(dir).unwrap();
924 }
925
926 #[test]
927 fn sha256_file_nonexistent_file_errors() {
928 assert!(sha256_file(Path::new("./no_such_file_for_sha256_file_test.bin")).is_err());
929 }
930
931 #[test]
932 fn read_latin1_decodes_every_byte() {
933 let dir = "./test_read_latin1_decodes_every_byte";
934 create_dir_all(dir).unwrap();
935 let path = Path::new(dir).join("latin1.txt");
936 // All 256 byte values, including sequences that are invalid UTF-8.
937 let bytes: Vec<u8> = (0u8..=255).collect();
938 fs::write(&path, &bytes).unwrap();
939
940 let text = read_latin1(&path).unwrap();
941
942 // One character per source byte, each with the byte's own scalar value.
943 assert_eq!(text.chars().count(), 256);
944 for (i, c) in text.chars().enumerate() {
945 assert_eq!(c as u32, i as u32, "byte {i} decoded to {c:?}");
946 }
947
948 remove_dir_all(dir).unwrap();
949 }
950
951 #[test]
952 fn read_latin1_empty_file_is_empty_string() {
953 let dir = "./test_read_latin1_empty_file_is_empty_string";
954 create_dir_all(dir).unwrap();
955 let path = Path::new(dir).join("empty.txt");
956 File::create(&path).unwrap();
957
958 assert_eq!(read_latin1(&path).unwrap(), "");
959
960 remove_dir_all(dir).unwrap();
961 }
962
963 #[test]
964 fn read_latin1_nonexistent_file_errors() {
965 assert!(read_latin1(Path::new("./no_such_file_for_read_latin1_test.txt")).is_err());
966 }
967
968 #[test]
969 fn download_to_with_retries_does_not_retry_unusable_url() {
970 let dir = "./test_download_to_with_retries_does_not_retry_unusable_url";
971 create_dir_all(dir).unwrap();
972
973 // A URL ending in `/` yields no filename. A retry cannot fix that, so this must
974 // fail immediately instead of sleeping through the retry schedule.
975 let started = std::time::Instant::now();
976 let result = download_to_with_retries("https://x.test/a/", Path::new(dir), None, 5);
977
978 assert!(matches!(result, Err(DatasetError::ValidationError(_))));
979 assert!(
980 started.elapsed() < RETRY_BASE_DELAY,
981 "a non-retryable error must not wait for a retry"
982 );
983
984 remove_dir_all(dir).unwrap();
985 }
986
987 #[test]
988 fn filename_from_url_plain_segment() {
989 assert_eq!(
990 filename_from_url("https://x.test/a/iris.csv"),
991 Some("iris.csv")
992 );
993 }
994
995 #[test]
996 fn filename_from_url_strips_query_and_fragment() {
997 assert_eq!(
998 filename_from_url("https://x.test/a/iris.csv?raw=1"),
999 Some("iris.csv")
1000 );
1001 assert_eq!(
1002 filename_from_url("https://x.test/a/iris.csv#section"),
1003 Some("iris.csv")
1004 );
1005 }
1006
1007 #[test]
1008 fn filename_from_url_trailing_slash_is_none() {
1009 assert_eq!(filename_from_url("https://x.test/a/"), None);
1010 }
1011
1012 #[test]
1013 fn filename_from_url_no_slash() {
1014 assert_eq!(filename_from_url("iris.csv"), Some("iris.csv"));
1015 }
1016
1017 #[test]
1018 fn inspect_cache_missing_then_fresh_and_stale() {
1019 let dir = "./test_inspect_cache_states";
1020 let dir_path = Path::new(dir);
1021 let dst = dir_path.join("data.txt");
1022 let _ = remove_dir_all(dir);
1023
1024 // Missing: the code creates the directory on demand. The file is absent.
1025 assert!(matches!(
1026 inspect_cache(dir_path, &dst, None).unwrap(),
1027 CacheState::Missing
1028 ));
1029 assert!(dir_path.exists());
1030
1031 // Fresh: file present, hash matches.
1032 fs::write(&dst, b"hello world").unwrap();
1033 assert!(matches!(
1034 inspect_cache(dir_path, &dst, Some(HELLO_WORLD_SHA256)).unwrap(),
1035 CacheState::Fresh
1036 ));
1037 // Fresh: no hash requested.
1038 assert!(matches!(
1039 inspect_cache(dir_path, &dst, None).unwrap(),
1040 CacheState::Fresh
1041 ));
1042
1043 // Stale: file present but hash mismatches.
1044 assert!(matches!(
1045 inspect_cache(dir_path, &dst, Some(ZERO_SHA256)).unwrap(),
1046 CacheState::Stale
1047 ));
1048
1049 remove_dir_all(dir).unwrap();
1050 }
1051}