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