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 # Parameters that the caller may set to override the defaults
  10 tempdir = None
  11 template = None
  12
  13 def gettempdir():
  14     """Function to calculate the directory to use."""
  15     global tempdir
  16     if tempdir is not None:
  17         return tempdir
  18     try:
  19         pwd = os.getcwd()
  20     except (AttributeError, os.error):
  21         pwd = os.curdir
  22     attempdirs = ['/var/tmp', '/usr/tmp', '/tmp', pwd]
  23     if os.name == 'nt':
  24         attempdirs.insert(0, 'C:\\TEMP')
  25         attempdirs.insert(0, '\\TEMP')
  26     elif os.name == 'mac':
  27         import macfs, MACFS
  28         try:
  29             refnum, dirid = macfs.FindFolder(MACFS.kOnSystemDisk,
  30                                              MACFS.kTemporaryFolderType, 1)
  31             dirname = macfs.FSSpec((refnum, dirid, '')).as_pathname()
  32             attempdirs.insert(0, dirname)
  33         except macfs.error:
  34             pass
  35     for envname in 'TMPDIR', 'TEMP', 'TMP':
  36         if os.environ.has_key(envname):
  37             attempdirs.insert(0, os.environ[envname])
  38     testfile = gettempprefix() + 'test'
  39     for dir in attempdirs:
  40         try:
  41             filename = os.path.join(dir, testfile)
  42             if os.name == 'posix':
  43                 try:
  44                     fd = os.open(filename,
  45                                  os.O_RDWR | os.O_CREAT | os.O_EXCL, 0700)
  46                 except OSError:
  47                     pass
  48                 else:
  49                     fp = os.fdopen(fd, 'w')
  50                     fp.write('blat')
  51                     fp.close()
  52                     os.unlink(filename)
  53                     del fp, fd
  54                     tempdir = dir
  55                     break
  56             else:
  57                 fp = open(filename, 'w')
  58                 fp.write('blat')
  59                 fp.close()
  60                 os.unlink(filename)
  61                 tempdir = dir
  62                 break
  63         except IOError:
  64             pass
  65     if tempdir is None:
  66         msg = "Can't find a usable temporary directory amongst " + `attempdirs`
  67         raise IOError, msg
  68     return tempdir
  69
  70
  71 # template caches the result of gettempprefix, for speed, when possible.
  72 # XXX unclear why this isn't "_template"; left it "template" for backward
  73 # compatibility.
  74 if os.name == "posix":
  75     # We don't try to cache the template on posix:  the pid may change on us
  76     # between calls due to a fork, and on Linux the pid changes even for
  77     # another thread in the same process.  Since any attempt to keep the
  78     # cache in synch would have to call os.getpid() anyway in order to make
  79     # sure the pid hasn't changed between calls, a cache wouldn't save any
  80     # time.  In addition, a cache is difficult to keep correct with the pid
  81     # changing willy-nilly, and earlier attempts proved buggy (races).
  82     template = None
  83
  84 # Else the pid never changes, so gettempprefix always returns the same
  85 # string.
  86 elif os.name == "nt":
  87     template = '~' + `os.getpid()` + '-'
  88 elif os.name == 'mac':
  89     template = 'Python-Tmp-'
  90 else:
  91     template = 'tmp' # XXX might choose a better one
  92
  93 def gettempprefix():
  94     """Function to calculate a prefix of the filename to use.
  95
  96     This incorporates the current process id on systems that support such a
  97     notion, so that concurrent processes don't generate the same prefix.
  98     """
  99
 100     global template
 101     if template is None:
 102         return '@' + `os.getpid()` + '.'
 103     else:
 104         return template
 105
 106
 107 def mktemp(suffix=""):
 108     """User-callable function to return a unique temporary file name."""
 109     dir = gettempdir()
 110     pre = gettempprefix()
 111     while 1:
 112         i = _counter.get_next()
 113         file = os.path.join(dir, pre + str(i) + suffix)
 114         if not os.path.exists(file):
 115             return file
 116
 117
 118 class TemporaryFileWrapper:
 119     """Temporary file wrapper
 120
 121     This class provides a wrapper around files opened for temporary use.
 122     In particular, it seeks to automatically remove the file when it is
 123     no longer needed.
 124     """
 125     def __init__(self, file, path):
 126         self.file = file
 127         self.path = path
 128
 129     def close(self):
 130         self.file.close()
 131         os.unlink(self.path)
 132
 133     def __del__(self):
 134         try: self.close()
 135         except: pass
 136
 137     def __getattr__(self, name):
 138         file = self.__dict__['file']
 139         a = getattr(file, name)
 140         if type(a) != type(0):
 141             setattr(self, name, a)
 142         return a
 143
 144
 145 def TemporaryFile(mode='w+b', bufsize=-1, suffix=""):
 146     """Create and return a temporary file (opened read-write by default)."""
 147     name = mktemp(suffix)
 148     if os.name == 'posix':
 149         # Unix -- be very careful
 150         fd = os.open(name, os.O_RDWR|os.O_CREAT|os.O_EXCL, 0700)
 151         try:
 152             os.unlink(name)
 153             return os.fdopen(fd, mode, bufsize)
 154         except:
 155             os.close(fd)
 156             raise
 157     else:
 158         # Non-unix -- can't unlink file that's still open, use wrapper
 159         file = open(name, mode, bufsize)
 160         return TemporaryFileWrapper(file, name)
 161
 162 # In order to generate unique names, mktemp() uses _counter.get_next().
 163 # This returns a unique integer on each call, in a threadsafe way (i.e.,
 164 # multiple threads will never see the same integer).  The integer will
 165 # usually be a Python int, but if _counter.get_next() is called often
 166 # enough, it will become a Python long.
 167 # Note that the only name that survives this next block of code
 168 # is "_counter".
 169
 170 class _ThreadSafeCounter:
 171     def __init__(self, mutex, initialvalue=0):
 172         self.mutex = mutex
 173         self.i = initialvalue
 174
 175     def get_next(self):
 176         self.mutex.acquire()
 177         result = self.i
 178         try:
 179             newi = result + 1
 180         except OverflowError:
 181             newi = long(result) + 1
 182         self.i = newi
 183         self.mutex.release()
 184         return result
 185
 186 try:
 187     import thread
 188
 189 except ImportError:
 190     class _DummyMutex:
 191         def acquire(self):
 192             pass
 193
 194         release = acquire
 195
 196     _counter = _ThreadSafeCounter(_DummyMutex())
 197     del _DummyMutex
 198
 199 else:
 200     _counter = _ThreadSafeCounter(thread.allocate_lock())
 201     del thread
 202
 203 del _ThreadSafeCounter