Lib/tempfile.py

   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 io as _io
  33 import os as _os
  34 import errno as _errno
  35 from random import Random as _Random
  36
  37 try:
  38     import fcntl as _fcntl
  39 except ImportError:
  40     def _set_cloexec(fd):
  41         pass
  42 else:
  43     def _set_cloexec(fd):
  44         try:
  45             flags = _fcntl.fcntl(fd, _fcntl.F_GETFD, 0)
  46         except IOError:
  47             pass
  48         else:
  49             # flags read successfully, modify
  50             flags |= _fcntl.FD_CLOEXEC
  51             _fcntl.fcntl(fd, _fcntl.F_SETFD, flags)
  52
  53
  54 try:
  55     import _thread
  56 except ImportError:
  57     import _dummy_thread as _thread
  58 _allocate_lock = _thread.allocate_lock
  59
  60 _text_openflags = _os.O_RDWR | _os.O_CREAT | _os.O_EXCL
  61 if hasattr(_os, 'O_NOINHERIT'):
  62     _text_openflags |= _os.O_NOINHERIT
  63 if hasattr(_os, 'O_NOFOLLOW'):
  64     _text_openflags |= _os.O_NOFOLLOW
  65
  66 _bin_openflags = _text_openflags
  67 if hasattr(_os, 'O_BINARY'):
  68     _bin_openflags |= _os.O_BINARY
  69
  70 if hasattr(_os, 'TMP_MAX'):
  71     TMP_MAX = _os.TMP_MAX
  72 else:
  73     TMP_MAX = 10000
  74
  75 template = "tmp"
  76
  77 # Internal routines.
  78
  79 _once_lock = _allocate_lock()
  80
  81 if hasattr(_os, "lstat"):
  82     _stat = _os.lstat
  83 elif hasattr(_os, "stat"):
  84     _stat = _os.stat
  85 else:
  86     # Fallback.  All we need is something that raises os.error if the
  87     # file doesn't exist.
  88     def _stat(fn):
  89         try:
  90             f = open(fn)
  91         except IOError:
  92             raise _os.error
  93         f.close()
  94
  95 def _exists(fn):
  96     try:
  97         _stat(fn)
  98     except _os.error:
  99         return False
 100     else:
 101         return True
 102
 103 class _RandomNameSequence:
 104     """An instance of _RandomNameSequence generates an endless
 105     sequence of unpredictable strings which can safely be incorporated
 106     into file names.  Each string is six characters long.  Multiple
 107     threads can safely use the same instance at the same time.
 108
 109     _RandomNameSequence is an iterator."""
 110
 111     characters = ("abcdefghijklmnopqrstuvwxyz" +
 112                   "ABCDEFGHIJKLMNOPQRSTUVWXYZ" +
 113                   "0123456789_")
 114
 115     def __init__(self):
 116         self.mutex = _allocate_lock()
 117         self.rng = _Random()
 118         self.normcase = _os.path.normcase
 119
 120     def __iter__(self):
 121         return self
 122
 123     def __next__(self):
 124         m = self.mutex
 125         c = self.characters
 126         choose = self.rng.choice
 127
 128         m.acquire()
 129         try:
 130             letters = [choose(c) for dummy in "123456"]
 131         finally:
 132             m.release()
 133
 134         return self.normcase(''.join(letters))
 135
 136 def _candidate_tempdir_list():
 137     """Generate a list of candidate temporary directories which
 138     _get_default_tempdir will try."""
 139
 140     dirlist = []
 141
 142     # First, try the environment.
 143     for envname in 'TMPDIR', 'TEMP', 'TMP':
 144         dirname = _os.getenv(envname)
 145         if dirname: dirlist.append(dirname)
 146
 147     # Failing that, try OS-specific locations.
 148     if _os.name == 'nt':
 149         dirlist.extend([ r'c:\temp', r'c:\tmp', r'\temp', r'\tmp' ])
 150     else:
 151         dirlist.extend([ '/tmp', '/var/tmp', '/usr/tmp' ])
 152
 153     # As a last resort, the current directory.
 154     try:
 155         dirlist.append(_os.getcwd())
 156     except (AttributeError, _os.error):
 157         dirlist.append(_os.curdir)
 158
 159     return dirlist
 160
 161 def _get_default_tempdir():
 162     """Calculate the default directory to use for temporary files.
 163     This routine should be called exactly once.
 164
 165     We determine whether or not a candidate temp dir is usable by
 166     trying to create and write to a file in that directory.  If this
 167     is successful, the test file is deleted.  To prevent denial of
 168     service, the name of the test file must be randomized."""
 169
 170     namer = _RandomNameSequence()
 171     dirlist = _candidate_tempdir_list()
 172     flags = _text_openflags
 173
 174     for dir in dirlist:
 175         if dir != _os.curdir:
 176             dir = _os.path.normcase(_os.path.abspath(dir))
 177         # Try only a few names per directory.
 178         for seq in range(100):
 179             name = next(namer)
 180             filename = _os.path.join(dir, name)
 181             try:
 182                 fd = _os.open(filename, flags, 0o600)
 183                 fp = _io.open(fd, 'wb')
 184                 fp.write(b'blat')
 185                 fp.close()
 186                 _os.unlink(filename)
 187                 del fp, fd
 188                 return dir
 189             except (OSError, IOError) as e:
 190                 if e.args[0] != _errno.EEXIST:
 191                     break # no point trying more names in this directory
 192                 pass
 193     raise IOError(_errno.ENOENT,
 194                   "No usable temporary directory found in %s" % dirlist)
 195
 196 _name_sequence = None
 197
 198 def _get_candidate_names():
 199     """Common setup sequence for all user-callable interfaces."""
 200
 201     global _name_sequence
 202     if _name_sequence is None:
 203         _once_lock.acquire()
 204         try:
 205             if _name_sequence is None:
 206                 _name_sequence = _RandomNameSequence()
 207         finally:
 208             _once_lock.release()
 209     return _name_sequence
 210
 211
 212 def _mkstemp_inner(dir, pre, suf, flags):
 213     """Code common to mkstemp, TemporaryFile, and NamedTemporaryFile."""
 214
 215     names = _get_candidate_names()
 216
 217     for seq in range(TMP_MAX):
 218         name = next(names)
 219         file = _os.path.join(dir, pre + name + suf)
 220         try:
 221             fd = _os.open(file, flags, 0o600)
 222             _set_cloexec(fd)
 223             return (fd, _os.path.abspath(file))
 224         except OSError as e:
 225             if e.errno == _errno.EEXIST:
 226                 continue # try again
 227             raise
 228
 229     raise IOError(_errno.EEXIST, "No usable temporary file name found")
 230
 231
 232 # User visible interfaces.
 233
 234 def gettempprefix():
 235     """Accessor for tempdir.template."""
 236     return template
 237
 238 tempdir = None
 239
 240 def gettempdir():
 241     """Accessor for tempfile.tempdir."""
 242     global tempdir
 243     if tempdir is None:
 244         _once_lock.acquire()
 245         try:
 246             if tempdir is None:
 247                 tempdir = _get_default_tempdir()
 248         finally:
 249             _once_lock.release()
 250     return tempdir
 251
 252 def mkstemp(suffix="", prefix=template, dir=None, text=False):
 253     """User-callable function to create and return a unique temporary
 254     file.  The return value is a pair (fd, name) where fd is the
 255     file descriptor returned by os.open, and name is the filename.
 256
 257     If 'suffix' is specified, the file name will end with that suffix,
 258     otherwise there will be no suffix.
 259
 260     If 'prefix' is specified, the file name will begin with that prefix,
 261     otherwise a default prefix is used.
 262
 263     If 'dir' is specified, the file will be created in that directory,
 264     otherwise a default directory is used.
 265
 266     If 'text' is specified and true, the file is opened in text
 267     mode.  Else (the default) the file is opened in binary mode.  On
 268     some operating systems, this makes no difference.
 269
 270     The file is readable and writable only by the creating user ID.
 271     If the operating system uses permission bits to indicate whether a
 272     file is executable, the file is executable by no one. The file
 273     descriptor is not inherited by children of this process.
 274
 275     Caller is responsible for deleting the file when done with it.
 276     """
 277
 278     if dir is None:
 279         dir = gettempdir()
 280
 281     if text:
 282         flags = _text_openflags
 283     else:
 284         flags = _bin_openflags
 285
 286     return _mkstemp_inner(dir, prefix, suffix, flags)
 287
 288
 289 def mkdtemp(suffix="", prefix=template, dir=None):
 290     """User-callable function to create and return a unique temporary
 291     directory.  The return value is the pathname of the directory.
 292
 293     Arguments are as for mkstemp, except that the 'text' argument is
 294     not accepted.
 295
 296     The directory is readable, writable, and searchable only by the
 297     creating user.
 298
 299     Caller is responsible for deleting the directory when done with it.
 300     """
 301
 302     if dir is None:
 303         dir = gettempdir()
 304
 305     names = _get_candidate_names()
 306
 307     for seq in range(TMP_MAX):
 308         name = next(names)
 309         file = _os.path.join(dir, prefix + name + suffix)
 310         try:
 311             _os.mkdir(file, 0o700)
 312             return file
 313         except OSError as e:
 314             if e.errno == _errno.EEXIST:
 315                 continue # try again
 316             raise
 317
 318     raise IOError(_errno.EEXIST, "No usable temporary directory name found")
 319
 320 def mktemp(suffix="", prefix=template, dir=None):
 321     """User-callable function to return a unique temporary file name.  The
 322     file is not created.
 323
 324     Arguments are as for mkstemp, except that the 'text' argument is
 325     not accepted.
 326
 327     This function is unsafe and should not be used.  The file name
 328     refers to a file that did not exist at some point, but by the time
 329     you get around to creating it, someone else may have beaten you to
 330     the punch.
 331     """
 332
 333 ##    from warnings import warn as _warn
 334 ##    _warn("mktemp is a potential security risk to your program",
 335 ##          RuntimeWarning, stacklevel=2)
 336
 337     if dir is None:
 338         dir = gettempdir()
 339
 340     names = _get_candidate_names()
 341     for seq in range(TMP_MAX):
 342         name = next(names)
 343         file = _os.path.join(dir, prefix + name + suffix)
 344         if not _exists(file):
 345             return file
 346
 347     raise IOError(_errno.EEXIST, "No usable temporary filename found")
 348
 349
 350 class _TemporaryFileWrapper:
 351     """Temporary file wrapper
 352
 353     This class provides a wrapper around files opened for
 354     temporary use.  In particular, it seeks to automatically
 355     remove the file when it is no longer needed.
 356     """
 357
 358     def __init__(self, file, name, delete=True):
 359         self.file = file
 360         self.name = name
 361         self.close_called = False
 362         self.delete = delete
 363
 364     def __getattr__(self, name):
 365         # Attribute lookups are delegated to the underlying file
 366         # and cached for non-numeric results
 367         # (i.e. methods are cached, closed and friends are not)
 368         file = self.__dict__['file']
 369         a = getattr(file, name)
 370         if not isinstance(a, int):
 371             setattr(self, name, a)
 372         return a
 373
 374     # The underlying __enter__ method returns the wrong object
 375     # (self.file) so override it to return the wrapper
 376     def __enter__(self):
 377         self.file.__enter__()
 378         return self
 379
 380     # iter() doesn't use __getattr__ to find the __iter__ method
 381     def __iter__(self):
 382         return iter(self.file)
 383
 384     # NT provides delete-on-close as a primitive, so we don't need
 385     # the wrapper to do anything special.  We still use it so that
 386     # file.name is useful (i.e. not "(fdopen)") with NamedTemporaryFile.
 387     if _os.name != 'nt':
 388         # Cache the unlinker so we don't get spurious errors at
 389         # shutdown when the module-level "os" is None'd out.  Note
 390         # that this must be referenced as self.unlink, because the
 391         # name TemporaryFileWrapper may also get None'd out before
 392         # __del__ is called.
 393         unlink = _os.unlink
 394
 395         def close(self):
 396             if not self.close_called:
 397                 self.close_called = True
 398                 self.file.close()
 399                 if self.delete:
 400                     self.unlink(self.name)
 401
 402         def __del__(self):
 403             self.close()
 404
 405         # Need to trap __exit__ as well to ensure the file gets
 406         # deleted when used in a with statement
 407         def __exit__(self, exc, value, tb):
 408             result = self.file.__exit__(exc, value, tb)
 409             self.close()
 410             return result
 411
 412
 413 def NamedTemporaryFile(mode='w+b', buffering=-1, encoding=None,
 414                        newline=None, suffix="", prefix=template,
 415                        dir=None, delete=True):
 416     """Create and return a temporary file.
 417     Arguments:
 418     'prefix', 'suffix', 'dir' -- as for mkstemp.
 419     'mode' -- the mode argument to io.open (default "w+b").
 420     'buffering' -- the buffer size argument to io.open (default -1).
 421     'encoding' -- the encoding argument to io.open (default None)
 422     'newline' -- the newline argument to io.open (default None)
 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 = _io.open(fd, mode, buffering=buffering,
 446                     newline=newline, encoding=encoding)
 447
 448     return _TemporaryFileWrapper(file, name, delete)
 449
 450 if _os.name != 'posix' or _os.sys.platform == 'cygwin':
 451     # On non-POSIX and Cygwin systems, assume that we cannot unlink a file
 452     # while it is open.
 453     TemporaryFile = NamedTemporaryFile
 454
 455 else:
 456     def TemporaryFile(mode='w+b', buffering=-1, encoding=None,
 457                       newline=None, suffix="", prefix=template,
 458                       dir=None):
 459         """Create and return a temporary file.
 460         Arguments:
 461         'prefix', 'suffix', 'dir' -- as for mkstemp.
 462         'mode' -- the mode argument to io.open (default "w+b").
 463         'buffering' -- the buffer size argument to io.open (default -1).
 464         'encoding' -- the encoding argument to io.open (default None)
 465         'newline' -- the newline argument to io.open (default None)
 466         The file is created as mkstemp() would do it.
 467
 468         Returns an object with a file-like interface.  The file has no
 469         name, and will cease to exist when it is closed.
 470         """
 471
 472         if dir is None:
 473             dir = gettempdir()
 474
 475         if 'b' in mode:
 476             flags = _bin_openflags
 477         else:
 478             flags = _text_openflags
 479
 480         (fd, name) = _mkstemp_inner(dir, prefix, suffix, flags)
 481         try:
 482             _os.unlink(name)
 483             return _io.open(fd, mode, buffering=buffering,
 484                             newline=newline, encoding=encoding)
 485         except:
 486             _os.close(fd)
 487             raise
 488
 489 class SpooledTemporaryFile:
 490     """Temporary file wrapper, specialized to switch from
 491     StringIO to a real file when it exceeds a certain size or
 492     when a fileno is needed.
 493     """
 494     _rolled = False
 495
 496     def __init__(self, max_size=0, mode='w+b', buffering=-1,
 497                  encoding=None, newline=None,
 498                  suffix="", prefix=template, dir=None):
 499         if 'b' in mode:
 500             self._file = _io.BytesIO()
 501         else:
 502             # Setting newline="\n" avoids newline translation;
 503             # this is important because otherwise on Windows we'd
 504             # hget double newline translation upon rollover().
 505             self._file = _io.StringIO(newline="\n")
 506         self._max_size = max_size
 507         self._rolled = False
 508         self._TemporaryFileArgs = {'mode': mode, 'buffering': buffering,
 509                                    'suffix': suffix, 'prefix': prefix,
 510                                    'encoding': encoding, 'newline': newline,
 511                                    'dir': dir}
 512
 513     def _check(self, file):
 514         if self._rolled: return
 515         max_size = self._max_size
 516         if max_size and file.tell() > max_size:
 517             self.rollover()
 518
 519     def rollover(self):
 520         if self._rolled: return
 521         file = self._file
 522         newfile = self._file = TemporaryFile(**self._TemporaryFileArgs)
 523         del self._TemporaryFileArgs
 524
 525         newfile.write(file.getvalue())
 526         newfile.seek(file.tell(), 0)
 527
 528         self._rolled = True
 529
 530     # The method caching trick from NamedTemporaryFile
 531     # won't work here, because _file may change from a
 532     # _StringIO instance to a real file. So we list
 533     # all the methods directly.
 534
 535     # Context management protocol
 536     def __enter__(self):
 537         if self._file.closed:
 538             raise ValueError("Cannot enter context with closed file")
 539         return self
 540
 541     def __exit__(self, exc, value, tb):
 542         self._file.close()
 543
 544     # file protocol
 545     def __iter__(self):
 546         return self._file.__iter__()
 547
 548     def close(self):
 549         self._file.close()
 550
 551     @property
 552     def closed(self):
 553         return self._file.closed
 554
 555     @property
 556     def encoding(self):
 557         return self._file.encoding
 558
 559     def fileno(self):
 560         self.rollover()
 561         return self._file.fileno()
 562
 563     def flush(self):
 564         self._file.flush()
 565
 566     def isatty(self):
 567         return self._file.isatty()
 568
 569     @property
 570     def mode(self):
 571         return self._file.mode
 572
 573     @property
 574     def name(self):
 575         return self._file.name
 576
 577     @property
 578     def newlines(self):
 579         return self._file.newlines
 580
 581     def next(self):
 582         return self._file.next
 583
 584     def read(self, *args):
 585         return self._file.read(*args)
 586
 587     def readline(self, *args):
 588         return self._file.readline(*args)
 589
 590     def readlines(self, *args):
 591         return self._file.readlines(*args)
 592
 593     def seek(self, *args):
 594         self._file.seek(*args)
 595
 596     @property
 597     def softspace(self):
 598         return self._file.softspace
 599
 600     def tell(self):
 601         return self._file.tell()
 602
 603     def truncate(self):
 604         self._file.truncate()
 605
 606     def write(self, s):
 607         file = self._file
 608         rv = file.write(s)
 609         self._check(file)
 610         return rv
 611
 612     def writelines(self, iterable):
 613         file = self._file
 614         rv = file.writelines(iterable)
 615         self._check(file)
 616         return rv
 617
 618     def xreadlines(self, *args):
 619         return self._file.xreadlines(*args)