1 files changed, 71 insertions, 6 deletions

diff --git a/tempfile/src/lib.rs b/tempfile/src/lib.rs
index 0b9e39b..92f766b 100644
--- a/tempfile/src/lib.rs
+++ b/tempfile/src/lib.rs
@@ -2,20 +2,39 @@
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
-//! Simplified tempfile which doesn't depend on the `rand` crate, instead using
-//! /dev/urandom as a source of entropy
+//! Simplified tempfile which doesn't depend on the `rand` crate.
+//!
+//! # Example
+//!
+//! ```
+//! use std::io::Result;
+//! use std::path::{Path, PathBuf};
+//! use tempfile::TempDir;
+//!
+//! fn main() -> Result<()> {
+//!     let t = TempDir::new()?;
+//!     assert!(t.path().exists());
+//!
+//!     Ok(())
+//! }
+//! ```
 
 use libc::mkdtemp;
 use std::env;
 use std::ffi::CString;
 use std::fs;
 use std::io::{Error, ErrorKind, Result};
+use std::mem::ManuallyDrop;
 use std::path::{Path, PathBuf};
+use std::ptr;
 
 pub struct Builder {
     prefix: String,
 }
 
+
+// Note: we implement a builder because the protoc-rust crate uses this API from
+// crates.io's tempfile. Our code mostly uses TempDir::new directly.
 impl Builder {
     pub fn new() -> Self {
         Builder {
@@ -31,10 +50,9 @@ impl Builder {
         self
     }
 
-    /// Tries to make a tempdir inside of `env::temp_dir()` with a specified
-    /// prefix. The directory and it's content is destroyed when TempDir is
-    /// dropped.
-    /// If the directory can not be created, `Err` is returned.
+    /// Creates a new temporary directory under libc's preferred system
+    /// temporary directory. The new directory will be removed when the returned
+    /// handle of type `TempDir` is dropped.
     pub fn tempdir(&self) -> Result<TempDir> {
         // mkdtemp() requires the template to end in 6 X chars, which will be replaced
         // with random characters to make the path unique.
@@ -68,17 +86,43 @@ impl Builder {
     }
 }
 
+/// Temporary directory. The directory will be removed when this object is
+/// dropped.
 pub struct TempDir {
     path: PathBuf,
+
+    // When adding new fields to TempDir: note that anything with a Drop impl
+    // will need to be dropped explicitly via ptr::read inside TempDir::remove
+    // or else it gets leaked (memory safe but not ideal).
 }
 
 impl TempDir {
+    pub fn new() -> Result<Self> {
+        Builder::new().tempdir()
+    }
+
     /// Accesses the tempdir's [`Path`].
     ///
     /// [`Path`]: http://doc.rust-lang.org/std/path/struct.Path.html
     pub fn path(&self) -> &Path {
         self.path.as_ref()
     }
+
+    /// Removes the temporary directory.
+    ///
+    /// Calling this is optional as dropping a TempDir object will also remove
+    /// the directory. Calling remove explicitly allows for any resulting error
+    /// to be handled.
+    pub fn remove(self) -> Result<()> {
+        // Place self inside ManuallyDrop so its Drop impl doesn't run, but nor
+        // does the path inside get dropped. Then use ptr::read to take out the
+        // PathBuf so that it *does* get dropped correctly at the bottom of this
+        // function.
+        let dont_drop = ManuallyDrop::new(self);
+        let path: PathBuf = unsafe { ptr::read(&dont_drop.path) };
+
+        fs::remove_dir_all(path)
+    }
 }
 
 impl Drop for TempDir {
@@ -86,3 +130,24 @@ impl Drop for TempDir {
         let _ = fs::remove_dir_all(&self.path);
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::TempDir;
+
+    #[test]
+    fn create_dir() {
+        let t = TempDir::new().unwrap();
+        let path = t.path();
+        assert!(path.exists());
+        assert!(path.is_dir());
+    }
+
+    #[test]
+    fn remove_dir() {
+        let t = TempDir::new().unwrap();
+        let path = t.path().to_owned();
+        assert!(t.remove().is_ok());
+        assert!(!path.exists());
+    }
+}