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     "SpooledTemporaryFile",

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

    24     "mktemp",                              # deprecated unsafe interface

    25     "TMP_MAX", "gettempprefix",            # constants

    26     "tempdir", "gettempdir"

    27    ]

    28 

    29 

    30 # Imports.

    31 

    32 import os as _os

    33 import errno as _errno

    34 from random import Random as _Random

    35 

    36 try:

    37     from cStringIO import StringIO as _StringIO

    38 except ImportError:

    39     from StringIO import StringIO as _StringIO

    40 

    41 try:

    42     import fcntl as _fcntl

    43 except ImportError:

    44     def _set_cloexec(fd):

    45         pass

    46 else:

    47     def _set_cloexec(fd):

    48         try:

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

    50         except IOError:

    51             pass

    52         else:

    53             # flags read successfully, modify

    54             flags |= _fcntl.FD_CLOEXEC

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

    56 

    57 

    58 try:

    59     import thread as _thread

    60 except ImportError:

    61     import dummy_thread as _thread

    62 _allocate_lock = _thread.allocate_lock

    63 

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

    65 if hasattr(_os, 'O_NOINHERIT'):

    66     _text_openflags |= _os.O_NOINHERIT

    67 if hasattr(_os, 'O_NOFOLLOW'):

    68     _text_openflags |= _os.O_NOFOLLOW

    69 

    70 _bin_openflags = _text_openflags

    71 if hasattr(_os, 'O_BINARY'):

    72     _bin_openflags |= _os.O_BINARY

    73 

    74 if hasattr(_os, 'TMP_MAX'):

    75     TMP_MAX = _os.TMP_MAX

    76 else:

    77     TMP_MAX = 10000

    78 

    79 template = "tmp"

    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 == 'riscos':

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

   154         if dirname: dirlist.append(dirname)

   155     elif _os.name == 'nt':

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

   157     else:

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

   159 

   160     # As a last resort, the current directory.

   161     try:

   162         dirlist.append(_os.getcwd())

   163     except (AttributeError, _os.error):

   164         dirlist.append(_os.curdir)

   165 

   166     return dirlist

   167 

   168 def _get_default_tempdir():

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

   170     This routine should be called exactly once.

   171 

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

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

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

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

   176 

   177     namer = _RandomNameSequence()

   178     dirlist = _candidate_tempdir_list()

   179     flags = _text_openflags

   180 

   181     for dir in dirlist:

   182         if dir != _os.curdir:

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

   184         # Try only a few names per directory.

   185         for seq in xrange(100):

   186             name = namer.next()

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

   188             try:

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

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

   191                 fp.write('blat')

   192                 fp.close()

   193                 _os.unlink(filename)

   194                 del fp, fd

   195                 return dir

   196             except (OSError, IOError), e:

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

   198                     break # no point trying more names in this directory

   199                 pass

   200     raise IOError, (_errno.ENOENT,

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

   202 

   203 _name_sequence = None

   204 

   205 def _get_candidate_names():

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

   207 

   208     global _name_sequence

   209     if _name_sequence is None:

   210         _once_lock.acquire()

   211         try:

   212             if _name_sequence is None:

   213                 _name_sequence = _RandomNameSequence()

   214         finally:

   215             _once_lock.release()

   216     return _name_sequence

   217 

   218 

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

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

   221 

   222     names = _get_candidate_names()

   223 

   224     for seq in xrange(TMP_MAX):

   225         name = names.next()

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

   227         try:

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

   229             _set_cloexec(fd)

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

   231         except OSError, e:

   232             if e.errno == _errno.EEXIST:

   233                 continue # try again

   234             raise

   235 

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

   237 

   238 

   239 # User visible interfaces.

   240 

   241 def gettempprefix():

   242     """Accessor for tempdir.template."""

   243     return template

   244 

   245 tempdir = None

   246 

   247 def gettempdir():

   248     """Accessor for tempfile.tempdir."""

   249     global tempdir

   250     if tempdir is None:

   251         _once_lock.acquire()

   252         try:

   253             if tempdir is None:

   254                 tempdir = _get_default_tempdir()

   255         finally:

   256             _once_lock.release()

   257     return tempdir

   258 

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

   260     """User-callable function to create and return a unique temporary

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

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

   263 

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

   265     otherwise there will be no suffix.

   266 

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

   268     otherwise a default prefix is used.

   269 

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

   271     otherwise a default directory is used.

   272 

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

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

   275     some operating systems, this makes no difference.

   276 

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

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

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

   280     descriptor is not inherited by children of this process.

   281 

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

   283     """

   284 

   285     if dir is None:

   286         dir = gettempdir()

   287 

   288     if text:

   289         flags = _text_openflags

   290     else:

   291         flags = _bin_openflags

   292 

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

   294 

   295 

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

   297     """User-callable function to create and return a unique temporary

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

   299 

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

   301     not accepted.

   302 

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

   304     creating user.

   305 

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

   307     """

   308 

   309     if dir is None:

   310         dir = gettempdir()

   311 

   312     names = _get_candidate_names()

   313 

   314     for seq in xrange(TMP_MAX):

   315         name = names.next()

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

   317         try:

   318             _os.mkdir(file, 0700)

   319             return file

   320         except OSError, e:

   321             if e.errno == _errno.EEXIST:

   322                 continue # try again

   323             raise

   324 

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

   326 

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

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

   329     file is not created.

   330 

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

   332     not accepted.

   333 

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

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

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

   337     the punch.

   338     """

   339 

   340 ##    from warnings import warn as _warn

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

   342 ##          RuntimeWarning, stacklevel=2)

   343 

   344     if dir is None:

   345         dir = gettempdir()

   346 

   347     names = _get_candidate_names()

   348     for seq in xrange(TMP_MAX):

   349         name = names.next()

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

   351         if not _exists(file):

   352             return file

   353 

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

   355 

   356 

   357 class _TemporaryFileWrapper:

   358     """Temporary file wrapper

   359 

   360     This class provides a wrapper around files opened for

   361     temporary use.  In particular, it seeks to automatically

   362     remove the file when it is no longer needed.

   363     """

   364 

   365     def __init__(self, file, name, delete=True):

   366         self.file = file

   367         self.name = name

   368         self.close_called = False

   369         self.delete = delete

   370 

   371     def __getattr__(self, name):

   372         # Attribute lookups are delegated to the underlying file

   373         # and cached for non-numeric results

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

   375         file = self.__dict__['file']

   376         a = getattr(file, name)

   377         if not issubclass(type(a), type(0)):

   378             setattr(self, name, a)

   379         return a

   380 

   381     # The underlying __enter__ method returns the wrong object

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

   383     def __enter__(self):

   384         self.file.__enter__()

   385         return self

   386 

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

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

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

   390     if _os.name != 'nt':

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

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

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

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

   395         # __del__ is called.

   396         unlink = _os.unlink

   397 

   398         def close(self):

   399             if not self.close_called:

   400                 self.close_called = True

   401                 self.file.close()

   402                 if self.delete:

   403                     self.unlink(self.name)

   404 

   405         def __del__(self):

   406             self.close()

   407 

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

   409         # deleted when used in a with statement

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

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

   412             self.close()

   413             return result

   414 

   415 

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

   417                        prefix=template, dir=None, delete=True):

   418     """Create and return a temporary file.

   419     Arguments:

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

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

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

   423     'delete' -- whether the file is deleted on close (default True).

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

   425 

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

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

   428     when it is closed unless the 'delete' argument is set to False.

   429     """

   430 

   431     if dir is None:

   432         dir = gettempdir()

   433 

   434     if 'b' in mode:

   435         flags = _bin_openflags

   436     else:

   437         flags = _text_openflags

   438 

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

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

   441     if _os.name == 'nt' and delete:

   442         flags |= _os.O_TEMPORARY

   443 

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

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

   446     return _TemporaryFileWrapper(file, name, delete)

   447 

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

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

   450     # while it is open.

   451     TemporaryFile = NamedTemporaryFile

   452 

   453 else:

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

   455                       prefix=template, dir=None):

   456         """Create and return a temporary file.

   457         Arguments:

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

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

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

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

   462 

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

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

   465         """

   466 

   467         if dir is None:

   468             dir = gettempdir()

   469 

   470         if 'b' in mode:

   471             flags = _bin_openflags

   472         else:

   473             flags = _text_openflags

   474 

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

   476         try:

   477             _os.unlink(name)

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

   479         except:

   480             _os.close(fd)

   481             raise

   482 

   483 class SpooledTemporaryFile:

   484     """Temporary file wrapper, specialized to switch from

   485     StringIO to a real file when it exceeds a certain size or

   486     when a fileno is needed.

   487     """

   488     _rolled = False

   489 

   490     def __init__(self, max_size=0, mode='w+b', bufsize=-1,

   491                  suffix="", prefix=template, dir=None):

   492         self._file = _StringIO()

   493         self._max_size = max_size

   494         self._rolled = False

   495         self._TemporaryFileArgs = (mode, bufsize, suffix, prefix, dir)

   496 

   497     def _check(self, file):

   498         if self._rolled: return

   499         max_size = self._max_size

   500         if max_size and file.tell() > max_size:

   501             self.rollover()

   502 

   503     def rollover(self):

   504         if self._rolled: return

   505         file = self._file

   506         newfile = self._file = TemporaryFile(*self._TemporaryFileArgs)

   507         del self._TemporaryFileArgs

   508 

   509         newfile.write(file.getvalue())

   510         newfile.seek(file.tell(), 0)

   511 

   512         self._rolled = True

   513 

   514     # The method caching trick from NamedTemporaryFile

   515     # won't work here, because _file may change from a

   516     # _StringIO instance to a real file. So we list

   517     # all the methods directly.

   518 

   519     # Context management protocol

   520     def __enter__(self):

   521         if self._file.closed:

   522             raise ValueError("Cannot enter context with closed file")

   523         return self

   524 

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

   526         self._file.close()

   527 

   528     # file protocol

   529     def __iter__(self):

   530         return self._file.__iter__()

   531 

   532     def close(self):

   533         self._file.close()

   534 

   535     @property

   536     def closed(self):

   537         return self._file.closed

   538 

   539     @property

   540     def encoding(self):

   541         return self._file.encoding

   542 

   543     def fileno(self):

   544         self.rollover()

   545         return self._file.fileno()

   546 

   547     def flush(self):

   548         self._file.flush()

   549 

   550     def isatty(self):

   551         return self._file.isatty()

   552 

   553     @property

   554     def mode(self):

   555         return self._file.mode

   556 

   557     @property

   558     def name(self):

   559         return self._file.name

   560 

   561     @property

   562     def newlines(self):

   563         return self._file.newlines

   564 

   565     def next(self):

   566         return self._file.next

   567 

   568     def read(self, *args):

   569         return self._file.read(*args)

   570 

   571     def readline(self, *args):

   572         return self._file.readline(*args)

   573 

   574     def readlines(self, *args):

   575         return self._file.readlines(*args)

   576 

   577     def seek(self, *args):

   578         self._file.seek(*args)

   579 

   580     @property

   581     def softspace(self):

   582         return self._file.softspace

   583 

   584     def tell(self):

   585         return self._file.tell()

   586 

   587     def truncate(self):

   588         self._file.truncate()

   589 

   590     def write(self, s):

   591         file = self._file

   592         rv = file.write(s)

   593         self._check(file)

   594         return rv

   595 

   596     def writelines(self, iterable):

   597         file = self._file

   598         rv = file.writelines(iterable)

   599         self._check(file)

   600         return rv

   601 

   602     def xreadlines(self, *args):

   603         return self._file.xreadlines(*args)