lib/stud/temporary.rb in stud-0.0.14 vs lib/stud/temporary.rb in stud-0.0.15

@@ -1,10 +1,13 @@
 require "securerandom" # for uuid generation
+require "stud/with"
 require "fileutils"
 
 module Stud
   module Temporary
+    include Stud::With
+
     # Returns a string for a randomly-generated temporary path.
     #
     # This does not create any files.
     def pathname(prefix="")
       root = ENV["TMP"] || ENV["TMPDIR"] || ENV["TEMP"] || "/tmp"
@@ -17,11 +20,22 @@
     # given to File.new.
     #
     # If no file args are given, the default file mode is "w+"
     def file(prefix="", *args, &block)
       args << "w+" if args.empty?
-      return File.new(pathname(prefix), *args)
+      if block_given?
+        with(File.new(pathname(prefix), *args)) do |fd|
+          begin
+            block.call(fd)
+            fd.close
+          ensure
+            File.unlink(fd.path)
+          end
+        end
+      else
+        return File.new(pathname(prefix), *args)
+      end
     end
 
     # Make a temporary directory.
     #
     # If given a block, the directory path is given to the block.  WHen the
@@ -32,11 +46,14 @@
     def directory(prefix="", &block)
       path = pathname(prefix)
       Dir.mkdir(path)
 
       if block_given?
-        block.call(path)
-        FileUtils.rm_r(path)
+        begin
+          block.call(path)
+        ensure
+          FileUtils.rm_r(path)
+        end
       else
         return path
       end
     end
   end # module Temporary