Lib/tempfile.py

   1 # Temporary file name allocation
   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
   8 import os
   9
  10
  11 # Parameters that the caller may set to override the defaults
  12
  13 tempdir = None
  14 template = None
  15
  16
  17 # Function to calculate the directory to use
  18
  19 def gettempdir():
  20     global tempdir
  21     if tempdir is not None:
  22         return tempdir
  23     try:
  24         pwd = os.getcwd()
  25     except (AttributeError, os.error):
  26         pwd = os.curdir
  27     attempdirs = ['/usr/tmp', '/tmp', pwd]
  28     if os.name == 'nt':
  29         attempdirs.insert(0, 'C:\\TEMP')
  30         attempdirs.insert(0, '\\TEMP')
  31     elif os.name == 'mac':
  32         import macfs, MACFS
  33         try:
  34              refnum, dirid = macfs.FindFolder(MACFS.kOnSystemDisk,
  35                                               MACFS.kTemporaryFolderType, 1)
  36              dirname = macfs.FSSpec((refnum, dirid, '')).as_pathname()
  37              attempdirs.insert(0, dirname)
  38         except macfs.error:
  39             pass
  40     for envname in 'TMPDIR', 'TEMP', 'TMP':
  41         if os.environ.has_key(envname):
  42             attempdirs.insert(0, os.environ[envname])
  43     testfile = gettempprefix() + 'test'
  44     for dir in attempdirs:
  45         try:
  46             filename = os.path.join(dir, testfile)
  47             fp = open(filename, 'w')
  48             fp.write('blat')
  49             fp.close()
  50             os.unlink(filename)
  51             tempdir = dir
  52             break
  53         except IOError:
  54             pass
  55     if tempdir is None:
  56         msg = "Can't find a usable temporary directory amongst " + `attempdirs`
  57         raise IOError, msg
  58     return tempdir
  59
  60
  61 # Function to calculate a prefix of the filename to use
  62
  63 _pid = None
  64
  65 def gettempprefix():
  66     global template, _pid
  67     if os.name == 'posix' and _pid and _pid != os.getpid():
  68         # Our pid changed; we must have forked -- zap the template
  69         template = None
  70     if template is None:
  71         if os.name == 'posix':
  72             _pid = os.getpid()
  73             template = '@' + `_pid` + '.'
  74         elif os.name == 'nt':
  75             template = '~' + `os.getpid()` + '-'
  76         elif os.name == 'mac':
  77             template = 'Python-Tmp-'
  78         else:
  79             template = 'tmp' # XXX might choose a better one
  80     return template
  81
  82
  83 # Counter for generating unique names
  84
  85 counter = 0
  86
  87
  88 # User-callable function to return a unique temporary file name
  89
  90 def mktemp(suffix=""):
  91     global counter
  92     dir = gettempdir()
  93     pre = gettempprefix()
  94     while 1:
  95         counter = counter + 1
  96         file = os.path.join(dir, pre + `counter` + suffix)
  97         if not os.path.exists(file):
  98             return file
  99
 100
 101 class TemporaryFileWrapper:
 102     """Temporary file wrapper
 103
 104     This class provides a wrapper around files opened for temporary use.
 105     In particular, it seeks to automatically remove the file when it is
 106     no longer needed.
 107     """
 108     def __init__(self, file, path):
 109         self.file = file
 110         self.path = path
 111
 112     def close(self):
 113         self.file.close()
 114         os.unlink(self.path)
 115
 116     def __del__(self):
 117         try: self.close()
 118         except: pass
 119
 120     def __getattr__(self, name):
 121         file = self.__dict__['file']
 122         a = getattr(file, name)
 123         setattr(self, name, a)
 124         return a
 125
 126
 127 def TemporaryFile(mode='w+b', bufsize=-1, suffix=""):
 128     name = mktemp(suffix)
 129     if os.name == 'posix':
 130         # Unix -- be very careful
 131         fd = os.open(name, os.O_RDWR|os.O_CREAT|os.O_EXCL, 0700)
 132         try:
 133             os.unlink(name)
 134             return os.fdopen(fd, mode, bufsize)
 135         except:
 136             os.close(fd)
 137             raise
 138     else:
 139         # Non-unix -- can't unlink file that's still open, use wrapper
 140         file = open(name, mode, bufsize)
 141         return TemporaryFileWrapper(file, name)