cpython/Lib/tempfile.py

# Temporary file name allocation
#
# XXX This tries to be not UNIX specific, but I don't know beans about
# how to choose a temp directory or filename on MS-DOS or other
# systems so it may have to be changed...


import os


# Parameters that the caller may set to override the defaults

tempdir = None
template = None


# Function to calculate the directory to use

def gettempdir():
    global tempdir
    if tempdir is not None:
        return tempdir
    try:
        pwd = os.getcwd()
    except (AttributeError, os.error):
        pwd = os.curdir
    attempdirs = ['/usr/tmp', '/tmp', pwd]
    if os.name == 'nt':
        attempdirs.insert(0, 'C:\\TEMP')
        attempdirs.insert(0, '\\TEMP')
    elif os.name == 'mac':
        import macfs, MACFS
        try:
             refnum, dirid = macfs.FindFolder(MACFS.kOnSystemDisk,
                                              MACFS.kTemporaryFolderType, 1)
             dirname = macfs.FSSpec((refnum, dirid, '')).as_pathname()
             attempdirs.insert(0, dirname)
        except macfs.error:
            pass
    for envname in 'TMPDIR', 'TEMP', 'TMP':
        if os.environ.has_key(envname):
            attempdirs.insert(0, os.environ[envname])
    testfile = gettempprefix() + 'test'
    for dir in attempdirs:
        try:
            filename = os.path.join(dir, testfile)
            fp = open(filename, 'w')
            fp.write('blat')
            fp.close()
            os.unlink(filename)
            tempdir = dir
            break
        except IOError:
            pass
    if tempdir is None:
        msg = "Can't find a usable temporary directory amongst " + `attempdirs`
        raise IOError, msg
    return tempdir


# Function to calculate a prefix of the filename to use

_pid = None

def gettempprefix():
    global template, _pid
    if os.name == 'posix' and _pid and _pid != os.getpid():
        # Our pid changed; we must have forked -- zap the template
        template = None
    if template is None:
        if os.name == 'posix':
            _pid = os.getpid()
            template = '@' + `_pid` + '.'
        elif os.name == 'nt':
            template = '~' + `os.getpid()` + '-'
        elif os.name == 'mac':
            template = 'Python-Tmp-'
        else:
            template = 'tmp' # XXX might choose a better one
    return template


# Counter for generating unique names

counter = 0


# User-callable function to return a unique temporary file name

def mktemp(suffix=""):
    global counter
    dir = gettempdir()
    pre = gettempprefix()
    while 1:
        counter = counter + 1
        file = os.path.join(dir, pre + `counter` + suffix)
        if not os.path.exists(file):
            return file


class TemporaryFileWrapper:
    """Temporary file wrapper

    This class provides a wrapper around files opened for temporary use.
    In particular, it seeks to automatically remove the file when it is
    no longer needed.
    """
    def __init__(self, file, path):
        self.file = file
        self.path = path

    def close(self):
        self.file.close()
        os.unlink(self.path)

    def __del__(self):
        try: self.close()
        except: pass

    def __getattr__(self, name):
        file = self.__dict__['file']
        a = getattr(file, name)
        setattr(self, name, a)
        return a


def TemporaryFile(mode='w+b', bufsize=-1, suffix=""):
    name = mktemp(suffix)
    if os.name == 'posix':
        # Unix -- be very careful
        fd = os.open(name, os.O_RDWR|os.O_CREAT|os.O_EXCL, 0700)
        try:
            os.unlink(name)
            return os.fdopen(fd, mode, bufsize)
        except:
            os.close(fd)
            raise
    else:
        # Non-unix -- can't unlink file that's still open, use wrapper
        file = open(name, mode, bufsize)
        return TemporaryFileWrapper(file, name)