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 def gettempprefix():
  64     global template
  65     if template == None:
  66         if os.name == 'posix':
  67             template = '@' + `os.getpid()` + '.'
  68         elif os.name == 'nt':
  69             template = '~' + `os.getpid()` + '-'
  70         elif os.name == 'mac':
  71             template = 'Python-Tmp-'
  72         else:
  73             template = 'tmp' # XXX might choose a better one
  74     return template
  75
  76
  77 # Counter for generating unique names
  78
  79 counter = 0
  80
  81
  82 # User-callable function to return a unique temporary file name
  83
  84 def mktemp(suffix=""):
  85     global counter
  86     dir = gettempdir()
  87     pre = gettempprefix()
  88     while 1:
  89         counter = counter + 1
  90         file = os.path.join(dir, pre + `counter` + suffix)
  91         if not os.path.exists(file):
  92             return file
  93
  94
  95 class TemporaryFileWrapper:
  96     """Temporary file wrapper
  97
  98     This class provides a wrapper around files opened for temporary use.
  99     In particular, it seeks to automatically remove the file when it is
 100     no longer needed.
 101     """
 102     def __init__(self, file, path):
 103         self.file = file
 104         self.path = path
 105
 106     def close(self):
 107         self.file.close()
 108         os.unlink(self.path)
 109
 110     def __del__(self):
 111         try: self.close()
 112         except: pass
 113
 114     def __getattr__(self, name):
 115         file = self.__dict__['file']
 116         a = getattr(file, name)
 117         setattr(self, name, a)
 118         return a
 119
 120
 121 def TemporaryFile(mode='w+b', bufsize=-1, suffix=""):
 122     name = mktemp(suffix)
 123     file = open(name, mode, bufsize)
 124     try:
 125         os.unlink(name)
 126     except os.error:
 127         # Non-unix -- can't unlink file that's still open, use wrapper
 128         return TemporaryFileWrapper(file, name)
 129     else:
 130         return file