Skip to main content

wordchipper_disk_cache/
disk_cache.rs

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