Lib/tempfile.py

   1 """Temporary files and filenames."""
   2
   3 # XXX This tries to be not UNIX specific, but I don't know beans about
   4 # how to choose a temp directory or filename on MS-DOS or other
   5 # systems so it may have to be changed...
   6
   7 import os
   8
   9 __all__ = ["mktemp", "TemporaryFile", "tempdir", "gettempprefix"]
  10
  11 # Parameters that the caller may set to override the defaults
  12 tempdir = None
  13 template = None
  14
  15 def gettempdir():
  16     """Function to calculate the directory to use."""
  17     global tempdir
  18     if tempdir is not None:
  19         return tempdir
  20
  21     # _gettempdir_inner deduces whether a candidate temp dir is usable by
  22     # trying to create a file in it, and write to it.  If that succeeds,
  23     # great, it closes the file and unlinks it.  There's a race, though:
  24     # the *name* of the test file it tries is the same across all threads
  25     # under most OSes (Linux is an exception), and letting multiple threads
  26     # all try to open, write to, close, and unlink a single file can cause
  27     # a variety of bogus errors (e.g., you cannot unlink a file under
  28     # Windows if anyone has it open, and two threads cannot create the
  29     # same file in O_EXCL mode under Unix).  The simplest cure is to serialize
  30     # calls to _gettempdir_inner.  This isn't a real expense, because the
  31     # first thread to succeed sets the global tempdir, and all subsequent
  32     # calls to gettempdir() reuse that without trying _gettempdir_inner.
  33     _tempdir_lock.acquire()
  34     try:
  35         return _gettempdir_inner()
  36     finally:
  37         _tempdir_lock.release()
  38
  39 def _gettempdir_inner():
  40     """Function to calculate the directory to use."""
  41     global tempdir
  42     if tempdir is not None:
  43         return tempdir
  44     try:
  45         pwd = os.getcwd()
  46     except (AttributeError, os.error):
  47         pwd = os.curdir
  48     attempdirs = ['/tmp', '/var/tmp', '/usr/tmp', pwd]
  49     if os.name == 'nt':
  50         attempdirs.insert(0, 'C:\\TEMP')
  51         attempdirs.insert(0, '\\TEMP')
  52     elif os.name == 'mac':
  53         import macfs, MACFS
  54         try:
  55             refnum, dirid = macfs.FindFolder(MACFS.kOnSystemDisk,
  56                                              MACFS.kTemporaryFolderType, 1)
  57             dirname = macfs.FSSpec((refnum, dirid, '')).as_pathname()
  58             attempdirs.insert(0, dirname)
  59         except macfs.error:
  60             pass
  61     elif os.name == 'riscos':
  62         scrapdir = os.getenv('Wimp$ScrapDir')
  63         if scrapdir:
  64             attempdirs.insert(0, scrapdir)
  65     for envname in 'TMPDIR', 'TEMP', 'TMP':
  66         if os.environ.has_key(envname):
  67             attempdirs.insert(0, os.environ[envname])
  68     testfile = gettempprefix() + 'test'
  69     for dir in attempdirs:
  70         try:
  71             filename = os.path.join(dir, testfile)
  72             if os.name == 'posix':
  73                 try:
  74                     fd = os.open(filename,
  75                                  os.O_RDWR | os.O_CREAT | os.O_EXCL, 0700)
  76                 except OSError:
  77                     pass
  78                 else:
  79                     fp = os.fdopen(fd, 'w')
  80                     fp.write('blat')
  81                     fp.close()
  82                     os.unlink(filename)
  83                     del fp, fd
  84                     tempdir = dir
  85                     break
  86             else:
  87                 fp = open(filename, 'w')
  88                 fp.write('blat')
  89                 fp.close()
  90                 os.unlink(filename)
  91                 tempdir = dir
  92                 break
  93         except IOError:
  94             pass
  95     if tempdir is None:
  96         msg = "Can't find a usable temporary directory amongst " + `attempdirs`
  97         raise IOError, msg
  98     return tempdir
  99
 100
 101 # template caches the result of gettempprefix, for speed, when possible.
 102 # XXX unclear why this isn't "_template"; left it "template" for backward
 103 # compatibility.
 104 if os.name == "posix":
 105     # We don't try to cache the template on posix:  the pid may change on us
 106     # between calls due to a fork, and on Linux the pid changes even for
 107     # another thread in the same process.  Since any attempt to keep the
 108     # cache in synch would have to call os.getpid() anyway in order to make
 109     # sure the pid hasn't changed between calls, a cache wouldn't save any
 110     # time.  In addition, a cache is difficult to keep correct with the pid
 111     # changing willy-nilly, and earlier attempts proved buggy (races).
 112     template = None
 113
 114 # Else the pid never changes, so gettempprefix always returns the same
 115 # string.
 116 elif os.name == "nt":
 117     template = '~' + `os.getpid()` + '-'
 118 elif os.name in ('mac', 'riscos'):
 119     template = 'Python-Tmp-'
 120 else:
 121     template = 'tmp' # XXX might choose a better one
 122
 123 def gettempprefix():
 124     """Function to calculate a prefix of the filename to use.
 125
 126     This incorporates the current process id on systems that support such a
 127     notion, so that concurrent processes don't generate the same prefix.
 128     """
 129
 130     global template
 131     if template is None:
 132         return '@' + `os.getpid()` + '.'
 133     else:
 134         return template
 135
 136
 137 def mktemp(suffix=""):
 138     """User-callable function to return a unique temporary file name."""
 139     dir = gettempdir()
 140     pre = gettempprefix()
 141     while 1:
 142         i = _counter.get_next()
 143         file = os.path.join(dir, pre + str(i) + suffix)
 144         if not os.path.exists(file):
 145             return file
 146
 147
 148 class TemporaryFileWrapper:
 149     """Temporary file wrapper
 150
 151     This class provides a wrapper around files opened for temporary use.
 152     In particular, it seeks to automatically remove the file when it is
 153     no longer needed.
 154     """
 155
 156     # Cache the unlinker so we don't get spurious errors at shutdown
 157     # when the module-level "os" is None'd out.  Note that this must
 158     # be referenced as self.unlink, because the name TemporaryFileWrapper
 159     # may also get None'd out before __del__ is called.
 160     unlink = os.unlink
 161
 162     def __init__(self, file, path):
 163         self.file = file
 164         self.path = path
 165         self.close_called = 0
 166
 167     def close(self):
 168         if not self.close_called:
 169             self.close_called = 1
 170             self.file.close()
 171             self.unlink(self.path)
 172
 173     def __del__(self):
 174         self.close()
 175
 176     def __getattr__(self, name):
 177         file = self.__dict__['file']
 178         a = getattr(file, name)
 179         if type(a) != type(0):
 180             setattr(self, name, a)
 181         return a
 182
 183
 184 def TemporaryFile(mode='w+b', bufsize=-1, suffix=""):
 185     """Create and return a temporary file (opened read-write by default)."""
 186     name = mktemp(suffix)
 187     if os.name == 'posix':
 188         # Unix -- be very careful
 189         fd = os.open(name, os.O_RDWR|os.O_CREAT|os.O_EXCL, 0700)
 190         try:
 191             os.unlink(name)
 192             return os.fdopen(fd, mode, bufsize)
 193         except:
 194             os.close(fd)
 195             raise
 196     else:
 197         # Non-unix -- can't unlink file that's still open, use wrapper
 198         file = open(name, mode, bufsize)
 199         return TemporaryFileWrapper(file, name)
 200
 201 # In order to generate unique names, mktemp() uses _counter.get_next().
 202 # This returns a unique integer on each call, in a threadsafe way (i.e.,
 203 # multiple threads will never see the same integer).  The integer will
 204 # usually be a Python int, but if _counter.get_next() is called often
 205 # enough, it will become a Python long.
 206 # Note that the only names that survive this next block of code
 207 # are "_counter" and "_tempdir_lock".
 208
 209 class _ThreadSafeCounter:
 210     def __init__(self, mutex, initialvalue=0):
 211         self.mutex = mutex
 212         self.i = initialvalue
 213
 214     def get_next(self):
 215         self.mutex.acquire()
 216         result = self.i
 217         try:
 218             newi = result + 1
 219         except OverflowError:
 220             newi = long(result) + 1
 221         self.i = newi
 222         self.mutex.release()
 223         return result
 224
 225 try:
 226     import thread
 227
 228 except ImportError:
 229     class _DummyMutex:
 230         def acquire(self):
 231             pass
 232
 233         release = acquire
 234
 235     _counter = _ThreadSafeCounter(_DummyMutex())
 236     _tempdir_lock = _DummyMutex()
 237     del _DummyMutex
 238
 239 else:
 240     _counter = _ThreadSafeCounter(thread.allocate_lock())
 241     _tempdir_lock = thread.allocate_lock()
 242     del thread
 243
 244 del _ThreadSafeCounter