Skip to main content

wordchipper_disk_cache/
disk_cache.rs

1//! # Wordchipper Disk Cache
2
3use std::{
4    fs,
5    path::{Path, PathBuf},
6};
7
8use anyhow::Context;
9use downloader::{Download, Downloader};
10
11use crate::{WORDCHIPPER_CACHE_CONFIG, path_utils};
12
13/// Options for [`WordchipperDiskCache`].
14#[derive(Clone, Default, Debug)]
15pub struct WordchipperDiskCacheOptions {
16    /// Optional path to the cache directory.
17    pub cache_dir: Option<PathBuf>,
18
19    /// Optional path to the data directory.
20    pub data_dir: Option<PathBuf>,
21
22    /// Optional [`Downloader`] builder.
23    pub downloader: Option<fn() -> Downloader>,
24}
25
26impl WordchipperDiskCacheOptions {
27    /// Set the cache directory.
28    pub fn with_cache_dir<P: AsRef<Path>>(
29        mut self,
30        cache_dir: Option<P>,
31    ) -> Self {
32        self.cache_dir = cache_dir.map(|p| p.as_ref().to_path_buf());
33        self
34    }
35
36    /// Set the data directory.
37    pub fn with_data_dir<P: AsRef<Path>>(
38        mut self,
39        data_dir: Option<P>,
40    ) -> Self {
41        self.data_dir = data_dir.map(|p| p.as_ref().to_path_buf());
42        self
43    }
44
45    /// Set the downloader builder.
46    pub fn with_downloader(
47        mut self,
48        downloader: Option<fn() -> Downloader>,
49    ) -> Self {
50        self.downloader = downloader;
51        self
52    }
53}
54
55/// Disk cache for downloaded files.
56///
57/// Leverages [`Downloader`] for downloading files,
58/// and [`PathResolver`](`crate::PathResolver`) for resolving cache and data paths
59/// appropriate for a user/system combo, and any environment overrides.
60pub struct WordchipperDiskCache {
61    /// Cache directory.
62    cache_dir: PathBuf,
63
64    /// Data directory.
65    data_dir: PathBuf,
66
67    /// Connection pool for downloading files.
68    downloader: Downloader,
69}
70
71impl Default for WordchipperDiskCache {
72    fn default() -> Self {
73        Self::new(WordchipperDiskCacheOptions::default()).unwrap()
74    }
75}
76
77impl WordchipperDiskCache {
78    /// Construct a new [`WordchipperDiskCache`].
79    pub fn new(options: WordchipperDiskCacheOptions) -> anyhow::Result<Self> {
80        let cache_dir = WORDCHIPPER_CACHE_CONFIG
81            .resolve_cache_dir(options.cache_dir)
82            .context("failed to resolve cache directory")?;
83
84        let data_dir = WORDCHIPPER_CACHE_CONFIG
85            .resolve_data_dir(options.data_dir)
86            .context("failed to resolve data directory")?;
87
88        let downloader = match options.downloader {
89            Some(builder) => builder(),
90            None => Downloader::builder().build()?,
91        };
92
93        Ok(Self {
94            cache_dir,
95            data_dir,
96            downloader,
97        })
98    }
99
100    /// Get the cache directory.
101    pub fn cache_dir(&self) -> &Path {
102        &self.cache_dir
103    }
104
105    /// Get the data directory.
106    pub fn data_dir(&self) -> &Path {
107        &self.data_dir
108    }
109
110    /// Get the downloader.
111    pub fn downloader(&self) -> &Downloader {
112        &self.downloader
113    }
114
115    /// Get the cache path for the given key.
116    ///
117    /// * Does not check that the path exists.
118    /// * Does not initialize the containing directories.
119    ///
120    /// # Arguments
121    /// * `context` - prefix dirs, inserted between `self.cache_dir` and `file`.
122    /// * `file` - the final file name.
123    pub fn cache_path<C, F>(
124        &self,
125        context: &[C],
126        file: F,
127    ) -> PathBuf
128    where
129        C: AsRef<Path>,
130        F: AsRef<Path>,
131    {
132        path_utils::extend_path(&self.cache_dir, context, file)
133    }
134
135    /// Loads a cached file from a specified path or downloads it if it does not exist.
136    ///
137    /// # Arguments
138    /// * `context`: A slice of `C` containing path-related context used in determining the
139    ///   cache location. These paths are combined to build the cached file's location.
140    /// * `urls`: A slice of string references specifying the URLs to download the file from
141    ///   if it is not already cached.
142    /// * `download`: A boolean flag indicating whether to attempt downloading the file
143    ///   from the provided URLs if it does not already exist in the cache.
144    ///
145    /// # Returns
146    /// * Returns a [`PathBuf`] pointing to the cached file if it exists or is successfully downloaded.
147    /// * Returns an error if the file is not found in the cache and downloading is not allowed
148    ///   or fails.
149    ///
150    /// # Errors
151    /// * Returns an error if the cached file does not exist and `download` is `false`.
152    /// * Returns an error if the downloading process fails.
153    pub fn load_cached_path<C, S>(
154        &mut self,
155        context: &[C],
156        urls: &[S],
157        download: bool,
158        /* TODO: hash: Option<&str>, */
159    ) -> anyhow::Result<PathBuf>
160    where
161        C: AsRef<Path>,
162        S: AsRef<str>,
163    {
164        let urls: Vec<_> = urls.iter().map(|s| s.as_ref()).collect();
165        let mut dl = Download::new_mirrored(&urls);
166        let file_name = dl.file_name.clone();
167        let path = self.cache_path(context, &file_name);
168        dl.file_name = path.clone();
169
170        if path.exists() {
171            return Ok(path);
172        }
173
174        if !download {
175            anyhow::bail!("cached file not found: {}", path.display());
176        }
177
178        fs::create_dir_all(path.parent().unwrap())?;
179
180        self.downloader.download(&[dl])?;
181
182        Ok(path)
183    }
184
185    /// Get the data path for the given key.
186    ///
187    /// * Does not check that the path exists.
188    /// * Does not initialize the containing directories.
189    ///
190    /// # Arguments
191    /// * `context` - prefix dirs, inserted between `self.cache_dir` and `file`.
192    /// * `file` - the final file name.
193    pub fn data_path<C, F>(
194        &self,
195        context: &[C],
196        file: F,
197    ) -> PathBuf
198    where
199        C: AsRef<Path>,
200        F: AsRef<Path>,
201    {
202        path_utils::extend_path(&self.data_dir, context, file)
203    }
204}
205
206#[cfg(test)]
207mod tests {
208    use std::{env, path::PathBuf};
209
210    use serial_test::serial;
211
212    use crate::{
213        WORDCHIPPER_CACHE_CONFIG,
214        WORDCHIPPER_CACHE_DIR,
215        WORDCHIPPER_DATA_DIR,
216        disk_cache::{WordchipperDiskCache, WordchipperDiskCacheOptions},
217    };
218
219    #[test]
220    #[serial]
221    fn test_resolve_dirs() {
222        let orig_cache_dir = env::var(WORDCHIPPER_CACHE_DIR);
223        let orig_data_dir = env::var(WORDCHIPPER_CACHE_DIR);
224
225        let pds = WORDCHIPPER_CACHE_CONFIG
226            .project_dirs()
227            .expect("failed to get project dirs");
228
229        let user_cache_dir = PathBuf::from("/tmp/wordchipper/cache");
230        let user_data_dir = PathBuf::from("/tmp/wordchipper/data");
231
232        let env_cache_dir = PathBuf::from("/tmp/wordchipper/env_cache");
233        let env_data_dir = PathBuf::from("/tmp/wordchipper/env_data");
234
235        // No env vars
236        unsafe {
237            env::remove_var(WORDCHIPPER_CACHE_DIR);
238            env::remove_var(WORDCHIPPER_DATA_DIR);
239        }
240
241        let cache = WordchipperDiskCache::new(
242            WordchipperDiskCacheOptions::default()
243                .with_cache_dir(Some(user_cache_dir.clone()))
244                .with_data_dir(Some(user_data_dir.clone())),
245        )
246        .unwrap();
247        assert_eq!(&cache.cache_dir(), &user_cache_dir);
248        assert_eq!(&cache.data_dir(), &user_data_dir);
249
250        let cache = WordchipperDiskCache::new(WordchipperDiskCacheOptions::default()).unwrap();
251        assert_eq!(&cache.cache_dir(), &pds.cache_dir().to_path_buf());
252        assert_eq!(&cache.data_dir(), &pds.data_dir().to_path_buf());
253
254        // With env var.
255        unsafe {
256            env::set_var(WORDCHIPPER_CACHE_DIR, env_cache_dir.to_str().unwrap());
257            env::set_var(WORDCHIPPER_DATA_DIR, env_data_dir.to_str().unwrap());
258        }
259
260        let cache = WordchipperDiskCache::new(
261            WordchipperDiskCacheOptions::default()
262                .with_cache_dir(Some(user_cache_dir.clone()))
263                .with_data_dir(Some(user_data_dir.clone())),
264        )
265        .unwrap();
266        assert_eq!(&cache.cache_dir(), &user_cache_dir);
267        assert_eq!(&cache.data_dir(), &user_data_dir);
268
269        let cache = WordchipperDiskCache::new(WordchipperDiskCacheOptions::default()).unwrap();
270        assert_eq!(&cache.cache_dir(), &env_cache_dir);
271        assert_eq!(&cache.data_dir(), &env_data_dir);
272
273        // restore original env var.
274        match orig_cache_dir {
275            Ok(original) => unsafe { env::set_var(WORDCHIPPER_CACHE_DIR, original) },
276            Err(_) => unsafe { env::remove_var(WORDCHIPPER_CACHE_DIR) },
277        }
278        match orig_data_dir {
279            Ok(original) => unsafe { env::set_var(WORDCHIPPER_DATA_DIR, original) },
280            Err(_) => unsafe { env::remove_var(WORDCHIPPER_DATA_DIR) },
281        }
282    }
283
284    #[test]
285    fn test_data_path() {
286        let cache = WordchipperDiskCache::new(WordchipperDiskCacheOptions::default()).unwrap();
287        let path = cache.data_path(&["prefix"], "file.txt");
288        assert_eq!(path, cache.data_dir.join("prefix").join("file.txt"));
289    }
290
291    #[test]
292    fn test_cache_path() {
293        let cache = WordchipperDiskCache::new(WordchipperDiskCacheOptions::default()).unwrap();
294        let path = cache.cache_path(&["prefix"], "file.txt");
295        assert_eq!(path, cache.cache_dir.join("prefix").join("file.txt"));
296    }
297}