1 """Temporary files.

     2 

     3 This module provides generic, low- and high-level interfaces for

     4 creating temporary files and directories.  The interfaces listed

     5 as "safe" just below can be used without fear of race conditions.

     6 Those listed as "unsafe" cannot, and are provided for backward

     7 compatibility only.

     8 

     9 This module also provides some data items to the user:

    10 

    11   TMP_MAX  - maximum number of names that will be tried before

    12              giving up.

    13   template - the default prefix for all temporary names.

    14              You may change this to control the default prefix.

    15   tempdir  - If this is set to a string before the first use of

    16              any routine from this module, it will be considered as

    17              another candidate location to store temporary files.

    18 """

    19 

    20 __all__ = [

    21     "NamedTemporaryFile", "TemporaryFile", # high level safe interfaces

    22     "mkstemp", "mkdtemp",                  # low level safe interfaces

    23     "mktemp",                              # deprecated unsafe interface

    24     "TMP_MAX", "gettempprefix",            # constants

    25     "tempdir", "gettempdir"

    26    ]

    27 

    28 

    29 # Imports.

    30 

    31 import os as _os

    32 import errno as _errno

    33 from random import Random as _Random

    34 

    35 if _os.name == 'mac':

    36     import Carbon.Folder as _Folder

    37     import Carbon.Folders as _Folders

    38 

    39 try:

    40     import fcntl as _fcntl

    41 except ImportError:

    42     def _set_cloexec(fd):

    43         pass

    44 else:

    45     def _set_cloexec(fd):

    46         try:

    47             flags = _fcntl.fcntl(fd, _fcntl.F_GETFD, 0)

    48         except IOError:

    49             pass

    50         else:

    51             # flags read successfully, modify

    52             flags |= _fcntl.FD_CLOEXEC

    53             _fcntl.fcntl(fd, _fcntl.F_SETFD, flags)

    54 

    55 

    56 try:

    57     import thread as _thread

    58 except ImportError:

    59     import dummy_thread as _thread

    60 _allocate_lock = _thread.allocate_lock

    61 

    62 _text_openflags = _os.O_RDWR | _os.O_CREAT | _os.O_EXCL

    63 if hasattr(_os, 'O_NOINHERIT'):

    64     _text_openflags |= _os.O_NOINHERIT

    65 if hasattr(_os, 'O_NOFOLLOW'):

    66     _text_openflags |= _os.O_NOFOLLOW

    67 

    68 _bin_openflags = _text_openflags

    69 if hasattr(_os, 'O_BINARY'):

    70     _bin_openflags |= _os.O_BINARY

    71 

    72 if hasattr(_os, 'TMP_MAX'):

    73     TMP_MAX = _os.TMP_MAX

    74 else:

    75     TMP_MAX = 10000

    76 

    77 template = "tmp"

    78 

    79 tempdir = None

    80 

    81 # Internal routines.

    82 

    83 _once_lock = _allocate_lock()

    84 

    85 if hasattr(_os, "lstat"):

    86     _stat = _os.lstat

    87 elif hasattr(_os, "stat"):

    88     _stat = _os.stat

    89 else:

    90     # Fallback.  All we need is something that raises os.error if the

    91     # file doesn't exist.

    92     def _stat(fn):

    93         try:

    94             f = open(fn)

    95         except IOError:

    96             raise _os.error

    97         f.close()

    98 

    99 def _exists(fn):

   100     try:

   101         _stat(fn)

   102     except _os.error:

   103         return False

   104     else:

   105         return True

   106 

   107 class _RandomNameSequence:

   108     """An instance of _RandomNameSequence generates an endless

   109     sequence of unpredictable strings which can safely be incorporated

   110     into file names.  Each string is six characters long.  Multiple

   111     threads can safely use the same instance at the same time.

   112 

   113     _RandomNameSequence is an iterator."""

   114 

   115     characters = ("abcdefghijklmnopqrstuvwxyz" +

   116                   "ABCDEFGHIJKLMNOPQRSTUVWXYZ" +

   117                   "0123456789-_")

   118 

   119     def __init__(self):

   120         self.mutex = _allocate_lock()

   121         self.rng = _Random()

   122         self.normcase = _os.path.normcase

   123 

   124     def __iter__(self):

   125         return self

   126 

   127     def next(self):

   128         m = self.mutex

   129         c = self.characters

   130         choose = self.rng.choice

   131 

   132         m.acquire()

   133         try:

   134             letters = [choose(c) for dummy in "123456"]

   135         finally:

   136             m.release()

   137 

   138         return self.normcase(''.join(letters))

   139 

   140 def _candidate_tempdir_list():

   141     """Generate a list of candidate temporary directories which

   142     _get_default_tempdir will try."""

   143 

   144     dirlist = []

   145 

   146     # First, try the environment.

   147     for envname in 'TMPDIR', 'TEMP', 'TMP':

   148         dirname = _os.getenv(envname)

   149         if dirname: dirlist.append(dirname)

   150 

   151     # Failing that, try OS-specific locations.

   152     if _os.name == 'mac':

   153         try:

   154             fsr = _Folder.FSFindFolder(_Folders.kOnSystemDisk,

   155                                               _Folders.kTemporaryFolderType, 1)

   156             dirname = fsr.as_pathname()

   157             dirlist.append(dirname)

   158         except _Folder.error:

   159             pass

   160     elif _os.name == 'riscos':

   161         dirname = _os.getenv('Wimp$ScrapDir')

   162         if dirname: dirlist.append(dirname)

   163     elif _os.name == 'nt':

   164         dirlist.extend([ r'c:\temp', r'c:\tmp', r'\temp', r'\tmp' ])

   165     else:

   166         dirlist.extend([ '/tmp', '/var/tmp', '/usr/tmp' ])

   167 

   168     # As a last resort, the current directory.

   169     try:

   170         dirlist.append(_os.getcwd())

   171     except (AttributeError, _os.error):

   172         dirlist.append(_os.curdir)

   173 

   174     return dirlist

   175 

   176 def _get_default_tempdir():

   177     """Calculate the default directory to use for temporary files.

   178     This routine should be called exactly once.

   179 

   180     We determine whether or not a candidate temp dir is usable by

   181     trying to create and write to a file in that directory.  If this

   182     is successful, the test file is deleted.  To prevent denial of

   183     service, the name of the test file must be randomized."""

   184 

   185     namer = _RandomNameSequence()

   186     dirlist = _candidate_tempdir_list()

   187     flags = _text_openflags

   188 

   189     for dir in dirlist:

   190         if dir != _os.curdir:

   191             dir = _os.path.normcase(_os.path.abspath(dir))

   192         # Try only a few names per directory.

   193         for seq in xrange(100):

   194             name = namer.next()

   195             filename = _os.path.join(dir, name)

   196             try:

   197                 fd = _os.open(filename, flags, 0600)

   198                 fp = _os.fdopen(fd, 'w')

   199                 fp.write('blat')

   200                 fp.close()

   201                 _os.unlink(filename)

   202                 del fp, fd

   203                 return dir

   204             except (OSError, IOError), e:

   205                 if e[0] != _errno.EEXIST:

   206                     break # no point trying more names in this directory

   207                 pass

   208     raise IOError, (_errno.ENOENT,

   209                     ("No usable temporary directory found in %s" % dirlist))

   210 

   211 _name_sequence = None

   212 

   213 def _get_candidate_names():

   214     """Common setup sequence for all user-callable interfaces."""

   215 

   216     global _name_sequence

   217     if _name_sequence is None:

   218         _once_lock.acquire()

   219         try:

   220             if _name_sequence is None:

   221                 _name_sequence = _RandomNameSequence()

   222         finally:

   223             _once_lock.release()

   224     return _name_sequence

   225 

   226 

   227 def _mkstemp_inner(dir, pre, suf, flags):

   228     """Code common to mkstemp, TemporaryFile, and NamedTemporaryFile."""

   229 

   230     names = _get_candidate_names()

   231 

   232     for seq in xrange(TMP_MAX):

   233         name = names.next()

   234         file = _os.path.join(dir, pre + name + suf)

   235         try:

   236             fd = _os.open(file, flags, 0600)

   237             _set_cloexec(fd)

   238             return (fd, _os.path.abspath(file))

   239         except OSError, e:

   240             if e.errno == _errno.EEXIST:

   241                 continue # try again

   242             raise

   243 

   244     raise IOError, (_errno.EEXIST, "No usable temporary file name found")

   245 

   246 

   247 # User visible interfaces.

   248 

   249 def gettempprefix():

   250     """Accessor for tempdir.template."""

   251     return template

   252 

   253 tempdir = None

   254 

   255 def gettempdir():

   256     """Accessor for tempdir.tempdir."""

   257     global tempdir

   258     if tempdir is None:

   259         _once_lock.acquire()

   260         try:

   261             if tempdir is None:

   262                 tempdir = _get_default_tempdir()

   263         finally:

   264             _once_lock.release()

   265     return tempdir

   266 

   267 def mkstemp(suffix="", prefix=template, dir=None, text=False):

   268     """mkstemp([suffix, [prefix, [dir, [text]]]])

   269     User-callable function to create and return a unique temporary

   270     file.  The return value is a pair (fd, name) where fd is the

   271     file descriptor returned by os.open, and name is the filename.

   272 

   273     If 'suffix' is specified, the file name will end with that suffix,

   274     otherwise there will be no suffix.

   275 

   276     If 'prefix' is specified, the file name will begin with that prefix,

   277     otherwise a default prefix is used.

   278 

   279     If 'dir' is specified, the file will be created in that directory,

   280     otherwise a default directory is used.

   281 

   282     If 'text' is specified and true, the file is opened in text

   283     mode.  Else (the default) the file is opened in binary mode.  On

   284     some operating systems, this makes no difference.

   285 

   286     The file is readable and writable only by the creating user ID.

   287     If the operating system uses permission bits to indicate whether a

   288     file is executable, the file is executable by no one. The file

   289     descriptor is not inherited by children of this process.

   290 

   291     Caller is responsible for deleting the file when done with it.

   292     """

   293 

   294     if dir is None:

   295         dir = gettempdir()

   296 

   297     if text:

   298         flags = _text_openflags

   299     else:

   300         flags = _bin_openflags

   301 

   302     return _mkstemp_inner(dir, prefix, suffix, flags)

   303 

   304 

   305 def mkdtemp(suffix="", prefix=template, dir=None):

   306     """mkdtemp([suffix, [prefix, [dir]]])

   307     User-callable function to create and return a unique temporary

   308     directory.  The return value is the pathname of the directory.

   309 

   310     Arguments are as for mkstemp, except that the 'text' argument is

   311     not accepted.

   312 

   313     The directory is readable, writable, and searchable only by the

   314     creating user.

   315 

   316     Caller is responsible for deleting the directory when done with it.

   317     """

   318 

   319     if dir is None:

   320         dir = gettempdir()

   321 

   322     names = _get_candidate_names()

   323 

   324     for seq in xrange(TMP_MAX):

   325         name = names.next()

   326         file = _os.path.join(dir, prefix + name + suffix)

   327         try:

   328             _os.mkdir(file, 0700)

   329             return file

   330         except OSError, e:

   331             if e.errno == _errno.EEXIST:

   332                 continue # try again

   333             raise

   334 

   335     raise IOError, (_errno.EEXIST, "No usable temporary directory name found")

   336 

   337 def mktemp(suffix="", prefix=template, dir=None):

   338     """mktemp([suffix, [prefix, [dir]]])

   339     User-callable function to return a unique temporary file name.  The

   340     file is not created.

   341 

   342     Arguments are as for mkstemp, except that the 'text' argument is

   343     not accepted.

   344 

   345     This function is unsafe and should not be used.  The file name

   346     refers to a file that did not exist at some point, but by the time

   347     you get around to creating it, someone else may have beaten you to

   348     the punch.

   349     """

   350 

   351 ##    from warnings import warn as _warn

   352 ##    _warn("mktemp is a potential security risk to your program",

   353 ##          RuntimeWarning, stacklevel=2)

   354 

   355     if dir is None:

   356         dir = gettempdir()

   357 

   358     names = _get_candidate_names()

   359     for seq in xrange(TMP_MAX):

   360         name = names.next()

   361         file = _os.path.join(dir, prefix + name + suffix)

   362         if not _exists(file):

   363             return file

   364 

   365     raise IOError, (_errno.EEXIST, "No usable temporary filename found")

   366 

   367 class _TemporaryFileWrapper:

   368     """Temporary file wrapper

   369 

   370     This class provides a wrapper around files opened for

   371     temporary use.  In particular, it seeks to automatically

   372     remove the file when it is no longer needed.

   373     """

   374 

   375     def __init__(self, file, name):

   376         self.file = file

   377         self.name = name

   378         self.close_called = False

   379 

   380     def __getattr__(self, name):

   381         # Attribute lookups are delegated to the underlying file

   382         # and cached for non-numeric results

   383         # (i.e. methods are cached, closed and friends are not)

   384         file = self.__dict__['file']

   385         a = getattr(file, name)

   386         if type(a) != type(0):

   387             setattr(self, name, a)

   388         return a

   389 

   390     # The underlying __enter__ method returns the wrong object

   391     # (self.file) so override it to return the wrapper

   392     def __enter__(self):

   393         self.file.__enter__()

   394         return self

   395 

   396     # NT provides delete-on-close as a primitive, so we don't need

   397     # the wrapper to do anything special.  We still use it so that

   398     # file.name is useful (i.e. not "(fdopen)") with NamedTemporaryFile.

   399     if _os.name != 'nt':

   400         # Cache the unlinker so we don't get spurious errors at

   401         # shutdown when the module-level "os" is None'd out.  Note

   402         # that this must be referenced as self.unlink, because the

   403         # name TemporaryFileWrapper may also get None'd out before

   404         # __del__ is called.

   405         unlink = _os.unlink

   406 

   407         def close(self):

   408             if not self.close_called:

   409                 self.close_called = True

   410                 self.file.close()

   411                 self.unlink(self.name)

   412 

   413         def __del__(self):

   414             self.close()

   415 

   416         # Need to trap __exit__ as well to ensure the file gets

   417         # deleted when used in a with statement

   418         def __exit__(self, exc, value, tb):

   419             result = self.file.__exit__(exc, value, tb)

   420             self.close()

   421             return result

   422 

   423 

   424 def NamedTemporaryFile(mode='w+b', bufsize=-1, suffix="",

   425                        prefix=template, dir=None):

   426     """Create and return a temporary file.

   427     Arguments:

   428     'prefix', 'suffix', 'dir' -- as for mkstemp.

   429     'mode' -- the mode argument to os.fdopen (default "w+b").

   430     'bufsize' -- the buffer size argument to os.fdopen (default -1).

   431     The file is created as mkstemp() would do it.

   432 

   433     Returns an object with a file-like interface; the name of the file

   434     is accessible as file.name.  The file will be automatically deleted

   435     when it is closed.

   436     """

   437 

   438     if dir is None:

   439         dir = gettempdir()

   440 

   441     if 'b' in mode:

   442         flags = _bin_openflags

   443     else:

   444         flags = _text_openflags

   445 

   446     # Setting O_TEMPORARY in the flags causes the OS to delete

   447     # the file when it is closed.  This is only supported by Windows.

   448     if _os.name == 'nt':

   449         flags |= _os.O_TEMPORARY

   450 

   451     (fd, name) = _mkstemp_inner(dir, prefix, suffix, flags)

   452     file = _os.fdopen(fd, mode, bufsize)

   453     return _TemporaryFileWrapper(file, name)

   454 

   455 if _os.name != 'posix' or _os.sys.platform == 'cygwin':

   456     # On non-POSIX and Cygwin systems, assume that we cannot unlink a file

   457     # while it is open.

   458     TemporaryFile = NamedTemporaryFile

   459 

   460 else:

   461     def TemporaryFile(mode='w+b', bufsize=-1, suffix="",

   462                       prefix=template, dir=None):

   463         """Create and return a temporary file.

   464         Arguments:

   465         'prefix', 'suffix', 'dir' -- as for mkstemp.

   466         'mode' -- the mode argument to os.fdopen (default "w+b").

   467         'bufsize' -- the buffer size argument to os.fdopen (default -1).

   468         The file is created as mkstemp() would do it.

   469 

   470         Returns an object with a file-like interface.  The file has no

   471         name, and will cease to exist when it is closed.

   472         """

   473 

   474         if dir is None:

   475             dir = gettempdir()

   476 

   477         if 'b' in mode:

   478             flags = _bin_openflags

   479         else:

   480             flags = _text_openflags

   481 

   482         (fd, name) = _mkstemp_inner(dir, prefix, suffix, flags)

   483         try:

   484             _os.unlink(name)

   485             return _os.fdopen(fd, mode, bufsize)

   486         except:

   487             _os.close(fd)

   488             raise