lib/tempfile.rb

   1 #
   2 # tempfile - manipulates temporary files
   3 #
   4 # $Id$
   5 #
   6
   7 require 'delegate'
   8 require 'tmpdir'
   9 require 'thread'
  10
  11 # A class for managing temporary files.  This library is written to be
  12 # thread safe.
  13 class Tempfile < DelegateClass(File)
  14   MAX_TRY = 10
  15   @@cleanlist = []
  16   @@lock = Mutex.new
  17
  18   # Creates a temporary file of mode 0600 in the temporary directory,
  19   # opens it with mode "w+", and returns a Tempfile object which
  20   # represents the created temporary file.  A Tempfile object can be
  21   # treated just like a normal File object.
  22   #
  23   # The basename parameter is used to determine the name of a
  24   # temporary file.  If an Array is given, the first element is used
  25   # as prefix string and the second as suffix string, respectively.
  26   # Otherwise it is treated as prefix string.
  27   #
  28   # If tmpdir is omitted, the temporary directory is determined by
  29   # Dir::tmpdir provided by 'tmpdir.rb'.
  30   # When $SAFE > 0 and the given tmpdir is tainted, it uses
  31   # /tmp. (Note that ENV values are tainted by default)
  32   def initialize(basename, tmpdir=Dir::tmpdir)
  33     if $SAFE > 0 and tmpdir.tainted?
  34       tmpdir = '/tmp'
  35     end
  36
  37     lock = tmpname = nil
  38     n = failure = 0
  39     @@lock.synchronize {
  40       begin
  41         begin
  42           tmpname = File.join(tmpdir, make_tmpname(basename, n))
  43           lock = tmpname + '.lock'
  44           n += 1
  45         end while @@cleanlist.include?(tmpname) or
  46             File.exist?(lock) or File.exist?(tmpname)
  47         Dir.mkdir(lock)
  48       rescue
  49         failure += 1
  50         retry if failure < MAX_TRY
  51         raise "cannot generate tempfile `%s'" % tmpname
  52       end
  53     }
  54
  55     @data = [tmpname]
  56     @clean_proc = Tempfile.callback(@data)
  57     ObjectSpace.define_finalizer(self, @clean_proc)
  58
  59     @tmpfile = File.open(tmpname, File::RDWR|File::CREAT|File::EXCL, 0600)
  60     @tmpname = tmpname
  61     @@cleanlist << @tmpname
  62     @data[1] = @tmpfile
  63     @data[2] = @@cleanlist
  64
  65     super(@tmpfile)
  66
  67     # Now we have all the File/IO methods defined, you must not
  68     # carelessly put bare puts(), etc. after this.
  69
  70     Dir.rmdir(lock)
  71   end
  72
  73   def make_tmpname(basename, n)
  74     case basename
  75     when Array
  76       prefix, suffix = *basename
  77     else
  78       prefix, suffix = basename, ''
  79     end
  80
  81     t = Time.now.strftime("%Y%m%d")
  82     path = "#{prefix}#{t}-#{$$}-#{rand(0x100000000).to_s(36)}-#{n}#{suffix}"
  83   end
  84   private :make_tmpname
  85
  86   # Opens or reopens the file with mode "r+".
  87   def open
  88     @tmpfile.close if @tmpfile
  89     @tmpfile = File.open(@tmpname, 'r+')
  90     @data[1] = @tmpfile
  91     __setobj__(@tmpfile)
  92   end
  93
  94   def _close    # :nodoc:
  95     @tmpfile.close if @tmpfile
  96     @tmpfile = nil
  97     @data[1] = nil if @data
  98   end
  99   protected :_close
 100
 101   #Closes the file.  If the optional flag is true, unlinks the file
 102   # after closing.
 103   #
 104   # If you don't explicitly unlink the temporary file, the removal
 105   # will be delayed until the object is finalized.
 106   def close(unlink_now=false)
 107     if unlink_now
 108       close!
 109     else
 110       _close
 111     end
 112   end
 113
 114   # Closes and unlinks the file.
 115   def close!
 116     _close
 117     @clean_proc.call
 118     ObjectSpace.undefine_finalizer(self)
 119     @data = @tmpname = nil
 120   end
 121
 122   # Unlinks the file.  On UNIX-like systems, it is often a good idea
 123   # to unlink a temporary file immediately after creating and opening
 124   # it, because it leaves other programs zero chance to access the
 125   # file.
 126   def unlink
 127     # keep this order for thread safeness
 128     begin
 129       File.unlink(@tmpname) if File.exist?(@tmpname)
 130       @@cleanlist.delete(@tmpname)
 131       @data = @tmpname = nil
 132       ObjectSpace.undefine_finalizer(self)
 133     rescue Errno::EACCES
 134       # may not be able to unlink on Windows; just ignore
 135     end
 136   end
 137   alias delete unlink
 138
 139   # Returns the full path name of the temporary file.
 140   def path
 141     @tmpname
 142   end
 143
 144   # Returns the size of the temporary file.  As a side effect, the IO
 145   # buffer is flushed before determining the size.
 146   def size
 147     if @tmpfile
 148       @tmpfile.flush
 149       @tmpfile.stat.size
 150     else
 151       0
 152     end
 153   end
 154   alias length size
 155
 156   class << self
 157     def callback(data)  # :nodoc:
 158       pid = $$
 159       Proc.new {
 160         if pid == $$
 161           path, tmpfile, cleanlist = *data
 162
 163           print "removing ", path, "..." if $DEBUG
 164
 165           tmpfile.close if tmpfile
 166
 167           # keep this order for thread safeness
 168           File.unlink(path) if File.exist?(path)
 169           cleanlist.delete(path) if cleanlist
 170
 171           print "done\n" if $DEBUG
 172         end
 173       }
 174     end
 175
 176     # If no block is given, this is a synonym for new().
 177     #
 178     # If a block is given, it will be passed tempfile as an argument,
 179     # and the tempfile will automatically be closed when the block
 180     # terminates.  In this case, open() returns nil.
 181     def open(*args)
 182       tempfile = new(*args)
 183
 184       if block_given?
 185         begin
 186           yield(tempfile)
 187         ensure
 188           tempfile.close
 189         end
 190
 191         nil
 192       else
 193         tempfile
 194       end
 195     end
 196   end
 197 end
 198
 199 if __FILE__ == $0
 200 #  $DEBUG = true
 201   f = Tempfile.new("foo")
 202   f.print("foo\n")
 203   f.close
 204   f.open
 205   p f.gets # => "foo\n"
 206   f.close!
 207 end