/Lib/tempfile.py
Python | 603 lines | 599 code | 0 blank | 4 comment | 1 complexity | 1eca671f8c81e26d7a78da7c62270c5d MD5 | raw file
1"""Temporary files. 2 3This module provides generic, low- and high-level interfaces for 4creating temporary files and directories. The interfaces listed 5as "safe" just below can be used without fear of race conditions. 6Those listed as "unsafe" cannot, and are provided for backward 7compatibility only. 8 9This 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 32import os as _os 33import errno as _errno 34from random import Random as _Random 35 36try: 37 from cStringIO import StringIO as _StringIO 38except ImportError: 39 from StringIO import StringIO as _StringIO 40 41try: 42 import fcntl as _fcntl 43except ImportError: 44 def _set_cloexec(fd): 45 pass 46else: 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 58try: 59 import thread as _thread 60except 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 65if hasattr(_os, 'O_NOINHERIT'): 66 _text_openflags |= _os.O_NOINHERIT 67if hasattr(_os, 'O_NOFOLLOW'): 68 _text_openflags |= _os.O_NOFOLLOW 69 70_bin_openflags = _text_openflags 71if hasattr(_os, 'O_BINARY'): 72 _bin_openflags |= _os.O_BINARY 73 74if hasattr(_os, 'TMP_MAX'): 75 TMP_MAX = _os.TMP_MAX 76else: 77 TMP_MAX = 10000 78 79template = "tmp" 80 81# Internal routines. 82 83_once_lock = _allocate_lock() 84 85if hasattr(_os, "lstat"): 86 _stat = _os.lstat 87elif hasattr(_os, "stat"): 88 _stat = _os.stat 89else: 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 99def _exists(fn): 100 try: 101 _stat(fn) 102 except _os.error: 103 return False 104 else: 105 return True 106 107class _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 140def _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 168def _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 205def _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 219def _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 241def gettempprefix(): 242 """Accessor for tempdir.template.""" 243 return template 244 245tempdir = None 246 247def 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 259def 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 296def 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 327def 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 357class _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 416def 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 448if _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 453else: 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 483class 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)