aptos_temppath/
lib.rs

1// Copyright (c) Aptos
2// SPDX-License-Identifier: Apache-2.0
3
4#![forbid(unsafe_code)]
5
6use rand::RngCore;
7use std::{
8    fs, io,
9    path::{Path, PathBuf},
10};
11
12/// A simple wrapper for creating a temporary directory that is automatically deleted when it's
13/// dropped. Used in lieu of tempfile due to the large number of dependencies.
14#[derive(Debug, PartialEq)]
15pub struct TempPath {
16    path_buf: PathBuf,
17    persist: bool,
18}
19
20impl Drop for TempPath {
21    fn drop(&mut self) {
22        if !self.persist {
23            fs::remove_dir_all(&self.path_buf)
24                .or_else(|_| fs::remove_file(&self.path_buf))
25                .unwrap_or(());
26        }
27    }
28}
29
30impl TempPath {
31    /// Create new, uninitialized temporary path in the system temp directory.
32    pub fn new() -> Self {
33        Self::new_with_temp_dir(std::env::temp_dir())
34    }
35
36    /// Create new, uninitialized temporary path in the specified directory.
37    pub fn new_with_temp_dir(temp_dir: PathBuf) -> Self {
38        let mut temppath = temp_dir;
39        let mut rng = rand::thread_rng();
40        let mut bytes = [0_u8; 16];
41        rng.fill_bytes(&mut bytes);
42        temppath.push(hex::encode(&bytes));
43
44        TempPath {
45            path_buf: temppath,
46            persist: false,
47        }
48    }
49
50    /// Return the underlying path to this temporary directory.
51    pub fn path(&self) -> &Path {
52        &self.path_buf
53    }
54
55    /// Keep the temp path
56    pub fn persist(&mut self) {
57        self.persist = true;
58    }
59
60    pub fn create_as_file(&self) -> io::Result<()> {
61        let mut builder = fs::OpenOptions::new();
62        builder.write(true).create_new(true);
63        builder.open(self.path())?;
64        Ok(())
65    }
66
67    pub fn create_as_dir(&self) -> io::Result<()> {
68        let builder = fs::DirBuilder::new();
69        builder.create(self.path())?;
70        Ok(())
71    }
72}
73
74impl std::convert::AsRef<Path> for TempPath {
75    fn as_ref(&self) -> &Path {
76        self.path()
77    }
78}
79
80impl Default for TempPath {
81    fn default() -> Self {
82        Self::new()
83    }
84}