tempfile.py 19 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640
  1. """Temporary files.
  2. This module provides generic, low- and high-level interfaces for
  3. creating temporary files and directories. All of the interfaces
  4. provided by this module can be used without fear of race conditions
  5. except for 'mktemp'. 'mktemp' is subject to race conditions and
  6. should not be used; it is provided for backward compatibility only.
  7. This module also provides some data items to the user:
  8. TMP_MAX - maximum number of names that will be tried before
  9. giving up.
  10. template - the default prefix for all temporary names.
  11. You may change this to control the default prefix.
  12. tempdir - If this is set to a string before the first use of
  13. any routine from this module, it will be considered as
  14. another candidate location to store temporary files.
  15. """
  16. __all__ = [
  17. "NamedTemporaryFile", "TemporaryFile", # high level safe interfaces
  18. "SpooledTemporaryFile",
  19. "mkstemp", "mkdtemp", # low level safe interfaces
  20. "mktemp", # deprecated unsafe interface
  21. "TMP_MAX", "gettempprefix", # constants
  22. "tempdir", "gettempdir"
  23. ]
  24. # Imports.
  25. import io as _io
  26. import os as _os
  27. import errno as _errno
  28. from random import Random as _Random
  29. try:
  30. from cStringIO import StringIO as _StringIO
  31. except ImportError:
  32. from StringIO import StringIO as _StringIO
  33. try:
  34. import fcntl as _fcntl
  35. except ImportError:
  36. def _set_cloexec(fd):
  37. pass
  38. else:
  39. def _set_cloexec(fd):
  40. try:
  41. flags = _fcntl.fcntl(fd, _fcntl.F_GETFD, 0)
  42. except IOError:
  43. pass
  44. else:
  45. # flags read successfully, modify
  46. flags |= _fcntl.FD_CLOEXEC
  47. _fcntl.fcntl(fd, _fcntl.F_SETFD, flags)
  48. try:
  49. import thread as _thread
  50. except ImportError:
  51. import dummy_thread as _thread
  52. _allocate_lock = _thread.allocate_lock
  53. _text_openflags = _os.O_RDWR | _os.O_CREAT | _os.O_EXCL
  54. if hasattr(_os, 'O_NOINHERIT'):
  55. _text_openflags |= _os.O_NOINHERIT
  56. if hasattr(_os, 'O_NOFOLLOW'):
  57. _text_openflags |= _os.O_NOFOLLOW
  58. _bin_openflags = _text_openflags
  59. if hasattr(_os, 'O_BINARY'):
  60. _bin_openflags |= _os.O_BINARY
  61. if hasattr(_os, 'TMP_MAX'):
  62. TMP_MAX = _os.TMP_MAX
  63. else:
  64. TMP_MAX = 10000
  65. template = "tmp"
  66. # Internal routines.
  67. _once_lock = _allocate_lock()
  68. if hasattr(_os, "lstat"):
  69. _stat = _os.lstat
  70. elif hasattr(_os, "stat"):
  71. _stat = _os.stat
  72. else:
  73. # Fallback. All we need is something that raises os.error if the
  74. # file doesn't exist.
  75. def _stat(fn):
  76. try:
  77. f = open(fn)
  78. except IOError:
  79. raise _os.error
  80. f.close()
  81. def _exists(fn):
  82. try:
  83. _stat(fn)
  84. except _os.error:
  85. return False
  86. else:
  87. return True
  88. class _RandomNameSequence:
  89. """An instance of _RandomNameSequence generates an endless
  90. sequence of unpredictable strings which can safely be incorporated
  91. into file names. Each string is six characters long. Multiple
  92. threads can safely use the same instance at the same time.
  93. _RandomNameSequence is an iterator."""
  94. characters = ("abcdefghijklmnopqrstuvwxyz" +
  95. "ABCDEFGHIJKLMNOPQRSTUVWXYZ" +
  96. "0123456789_")
  97. def __init__(self):
  98. self.mutex = _allocate_lock()
  99. self.normcase = _os.path.normcase
  100. @property
  101. def rng(self):
  102. cur_pid = _os.getpid()
  103. if cur_pid != getattr(self, '_rng_pid', None):
  104. self._rng = _Random()
  105. self._rng_pid = cur_pid
  106. return self._rng
  107. def __iter__(self):
  108. return self
  109. def next(self):
  110. m = self.mutex
  111. c = self.characters
  112. choose = self.rng.choice
  113. m.acquire()
  114. try:
  115. letters = [choose(c) for dummy in "123456"]
  116. finally:
  117. m.release()
  118. return self.normcase(''.join(letters))
  119. def _candidate_tempdir_list():
  120. """Generate a list of candidate temporary directories which
  121. _get_default_tempdir will try."""
  122. dirlist = []
  123. # First, try the environment.
  124. for envname in 'TMPDIR', 'TEMP', 'TMP':
  125. dirname = _os.getenv(envname)
  126. if dirname: dirlist.append(dirname)
  127. # Failing that, try OS-specific locations.
  128. if _os.name == 'riscos':
  129. dirname = _os.getenv('Wimp$ScrapDir')
  130. if dirname: dirlist.append(dirname)
  131. elif _os.name == 'nt':
  132. dirlist.extend([ r'c:\temp', r'c:\tmp', r'\temp', r'\tmp' ])
  133. else:
  134. dirlist.extend([ '/tmp', '/var/tmp', '/usr/tmp' ])
  135. # As a last resort, the current directory.
  136. try:
  137. dirlist.append(_os.getcwd())
  138. except (AttributeError, _os.error):
  139. dirlist.append(_os.curdir)
  140. return dirlist
  141. def _get_default_tempdir():
  142. """Calculate the default directory to use for temporary files.
  143. This routine should be called exactly once.
  144. We determine whether or not a candidate temp dir is usable by
  145. trying to create and write to a file in that directory. If this
  146. is successful, the test file is deleted. To prevent denial of
  147. service, the name of the test file must be randomized."""
  148. namer = _RandomNameSequence()
  149. dirlist = _candidate_tempdir_list()
  150. flags = _text_openflags
  151. for dir in dirlist:
  152. if dir != _os.curdir:
  153. dir = _os.path.normcase(_os.path.abspath(dir))
  154. # Try only a few names per directory.
  155. for seq in xrange(100):
  156. name = namer.next()
  157. filename = _os.path.join(dir, name)
  158. try:
  159. fd = _os.open(filename, flags, 0o600)
  160. try:
  161. try:
  162. with _io.open(fd, 'wb', closefd=False) as fp:
  163. fp.write(b'blat')
  164. finally:
  165. _os.close(fd)
  166. finally:
  167. _os.unlink(filename)
  168. return dir
  169. except (OSError, IOError) as e:
  170. if e.args[0] == _errno.EEXIST:
  171. continue
  172. if (_os.name == 'nt' and e.args[0] == _errno.EACCES and
  173. _os.path.isdir(dir) and _os.access(dir, _os.W_OK)):
  174. # On windows, when a directory with the chosen name already
  175. # exists, EACCES error code is returned instead of EEXIST.
  176. continue
  177. break # no point trying more names in this directory
  178. raise IOError, (_errno.ENOENT,
  179. ("No usable temporary directory found in %s" % dirlist))
  180. _name_sequence = None
  181. def _get_candidate_names():
  182. """Common setup sequence for all user-callable interfaces."""
  183. global _name_sequence
  184. if _name_sequence is None:
  185. _once_lock.acquire()
  186. try:
  187. if _name_sequence is None:
  188. _name_sequence = _RandomNameSequence()
  189. finally:
  190. _once_lock.release()
  191. return _name_sequence
  192. def _mkstemp_inner(dir, pre, suf, flags):
  193. """Code common to mkstemp, TemporaryFile, and NamedTemporaryFile."""
  194. names = _get_candidate_names()
  195. for seq in xrange(TMP_MAX):
  196. name = names.next()
  197. file = _os.path.join(dir, pre + name + suf)
  198. try:
  199. fd = _os.open(file, flags, 0600)
  200. _set_cloexec(fd)
  201. return (fd, _os.path.abspath(file))
  202. except OSError, e:
  203. if e.errno == _errno.EEXIST:
  204. continue # try again
  205. if (_os.name == 'nt' and e.errno == _errno.EACCES and
  206. _os.path.isdir(dir) and _os.access(dir, _os.W_OK)):
  207. # On windows, when a directory with the chosen name already
  208. # exists, EACCES error code is returned instead of EEXIST.
  209. continue
  210. raise
  211. raise IOError, (_errno.EEXIST, "No usable temporary file name found")
  212. # User visible interfaces.
  213. def gettempprefix():
  214. """Accessor for tempdir.template."""
  215. return template
  216. tempdir = None
  217. def gettempdir():
  218. """Accessor for tempfile.tempdir."""
  219. global tempdir
  220. if tempdir is None:
  221. _once_lock.acquire()
  222. try:
  223. if tempdir is None:
  224. tempdir = _get_default_tempdir()
  225. finally:
  226. _once_lock.release()
  227. return tempdir
  228. def mkstemp(suffix="", prefix=template, dir=None, text=False):
  229. """User-callable function to create and return a unique temporary
  230. file. The return value is a pair (fd, name) where fd is the
  231. file descriptor returned by os.open, and name is the filename.
  232. If 'suffix' is specified, the file name will end with that suffix,
  233. otherwise there will be no suffix.
  234. If 'prefix' is specified, the file name will begin with that prefix,
  235. otherwise a default prefix is used.
  236. If 'dir' is specified, the file will be created in that directory,
  237. otherwise a default directory is used.
  238. If 'text' is specified and true, the file is opened in text
  239. mode. Else (the default) the file is opened in binary mode. On
  240. some operating systems, this makes no difference.
  241. The file is readable and writable only by the creating user ID.
  242. If the operating system uses permission bits to indicate whether a
  243. file is executable, the file is executable by no one. The file
  244. descriptor is not inherited by children of this process.
  245. Caller is responsible for deleting the file when done with it.
  246. """
  247. if dir is None:
  248. dir = gettempdir()
  249. if text:
  250. flags = _text_openflags
  251. else:
  252. flags = _bin_openflags
  253. return _mkstemp_inner(dir, prefix, suffix, flags)
  254. def mkdtemp(suffix="", prefix=template, dir=None):
  255. """User-callable function to create and return a unique temporary
  256. directory. The return value is the pathname of the directory.
  257. Arguments are as for mkstemp, except that the 'text' argument is
  258. not accepted.
  259. The directory is readable, writable, and searchable only by the
  260. creating user.
  261. Caller is responsible for deleting the directory when done with it.
  262. """
  263. if dir is None:
  264. dir = gettempdir()
  265. names = _get_candidate_names()
  266. for seq in xrange(TMP_MAX):
  267. name = names.next()
  268. file = _os.path.join(dir, prefix + name + suffix)
  269. try:
  270. _os.mkdir(file, 0700)
  271. return file
  272. except OSError, e:
  273. if e.errno == _errno.EEXIST:
  274. continue # try again
  275. if (_os.name == 'nt' and e.errno == _errno.EACCES and
  276. _os.path.isdir(dir) and _os.access(dir, _os.W_OK)):
  277. # On windows, when a directory with the chosen name already
  278. # exists, EACCES error code is returned instead of EEXIST.
  279. continue
  280. raise
  281. raise IOError, (_errno.EEXIST, "No usable temporary directory name found")
  282. def mktemp(suffix="", prefix=template, dir=None):
  283. """User-callable function to return a unique temporary file name. The
  284. file is not created.
  285. Arguments are as for mkstemp, except that the 'text' argument is
  286. not accepted.
  287. This function is unsafe and should not be used. The file name
  288. refers to a file that did not exist at some point, but by the time
  289. you get around to creating it, someone else may have beaten you to
  290. the punch.
  291. """
  292. ## from warnings import warn as _warn
  293. ## _warn("mktemp is a potential security risk to your program",
  294. ## RuntimeWarning, stacklevel=2)
  295. if dir is None:
  296. dir = gettempdir()
  297. names = _get_candidate_names()
  298. for seq in xrange(TMP_MAX):
  299. name = names.next()
  300. file = _os.path.join(dir, prefix + name + suffix)
  301. if not _exists(file):
  302. return file
  303. raise IOError, (_errno.EEXIST, "No usable temporary filename found")
  304. class _TemporaryFileWrapper:
  305. """Temporary file wrapper
  306. This class provides a wrapper around files opened for
  307. temporary use. In particular, it seeks to automatically
  308. remove the file when it is no longer needed.
  309. """
  310. def __init__(self, file, name, delete=True):
  311. self.file = file
  312. self.name = name
  313. self.close_called = False
  314. self.delete = delete
  315. def __getattr__(self, name):
  316. # Attribute lookups are delegated to the underlying file
  317. # and cached for non-numeric results
  318. # (i.e. methods are cached, closed and friends are not)
  319. file = self.__dict__['file']
  320. a = getattr(file, name)
  321. if not issubclass(type(a), type(0)):
  322. setattr(self, name, a)
  323. return a
  324. # The underlying __enter__ method returns the wrong object
  325. # (self.file) so override it to return the wrapper
  326. def __enter__(self):
  327. self.file.__enter__()
  328. return self
  329. # NT provides delete-on-close as a primitive, so we don't need
  330. # the wrapper to do anything special. We still use it so that
  331. # file.name is useful (i.e. not "(fdopen)") with NamedTemporaryFile.
  332. if _os.name != 'nt':
  333. # Cache the unlinker so we don't get spurious errors at
  334. # shutdown when the module-level "os" is None'd out. Note
  335. # that this must be referenced as self.unlink, because the
  336. # name TemporaryFileWrapper may also get None'd out before
  337. # __del__ is called.
  338. unlink = _os.unlink
  339. def close(self):
  340. if not self.close_called:
  341. self.close_called = True
  342. try:
  343. self.file.close()
  344. finally:
  345. if self.delete:
  346. self.unlink(self.name)
  347. def __del__(self):
  348. self.close()
  349. # Need to trap __exit__ as well to ensure the file gets
  350. # deleted when used in a with statement
  351. def __exit__(self, exc, value, tb):
  352. result = self.file.__exit__(exc, value, tb)
  353. self.close()
  354. return result
  355. else:
  356. def __exit__(self, exc, value, tb):
  357. self.file.__exit__(exc, value, tb)
  358. def NamedTemporaryFile(mode='w+b', bufsize=-1, suffix="",
  359. prefix=template, dir=None, delete=True):
  360. """Create and return a temporary file.
  361. Arguments:
  362. 'prefix', 'suffix', 'dir' -- as for mkstemp.
  363. 'mode' -- the mode argument to os.fdopen (default "w+b").
  364. 'bufsize' -- the buffer size argument to os.fdopen (default -1).
  365. 'delete' -- whether the file is deleted on close (default True).
  366. The file is created as mkstemp() would do it.
  367. Returns an object with a file-like interface; the name of the file
  368. is accessible as its 'name' attribute. The file will be automatically
  369. deleted when it is closed unless the 'delete' argument is set to False.
  370. """
  371. if dir is None:
  372. dir = gettempdir()
  373. if 'b' in mode:
  374. flags = _bin_openflags
  375. else:
  376. flags = _text_openflags
  377. # Setting O_TEMPORARY in the flags causes the OS to delete
  378. # the file when it is closed. This is only supported by Windows.
  379. if _os.name == 'nt' and delete:
  380. flags |= _os.O_TEMPORARY
  381. (fd, name) = _mkstemp_inner(dir, prefix, suffix, flags)
  382. try:
  383. file = _os.fdopen(fd, mode, bufsize)
  384. return _TemporaryFileWrapper(file, name, delete)
  385. except BaseException:
  386. _os.unlink(name)
  387. _os.close(fd)
  388. raise
  389. if _os.name != 'posix' or _os.sys.platform == 'cygwin':
  390. # On non-POSIX and Cygwin systems, assume that we cannot unlink a file
  391. # while it is open.
  392. TemporaryFile = NamedTemporaryFile
  393. else:
  394. def TemporaryFile(mode='w+b', bufsize=-1, suffix="",
  395. prefix=template, dir=None):
  396. """Create and return a temporary file.
  397. Arguments:
  398. 'prefix', 'suffix', 'dir' -- as for mkstemp.
  399. 'mode' -- the mode argument to os.fdopen (default "w+b").
  400. 'bufsize' -- the buffer size argument to os.fdopen (default -1).
  401. The file is created as mkstemp() would do it.
  402. Returns an object with a file-like interface. The file has no
  403. name, and will cease to exist when it is closed.
  404. """
  405. if dir is None:
  406. dir = gettempdir()
  407. if 'b' in mode:
  408. flags = _bin_openflags
  409. else:
  410. flags = _text_openflags
  411. (fd, name) = _mkstemp_inner(dir, prefix, suffix, flags)
  412. try:
  413. _os.unlink(name)
  414. return _os.fdopen(fd, mode, bufsize)
  415. except:
  416. _os.close(fd)
  417. raise
  418. class SpooledTemporaryFile:
  419. """Temporary file wrapper, specialized to switch from
  420. StringIO to a real file when it exceeds a certain size or
  421. when a fileno is needed.
  422. """
  423. _rolled = False
  424. def __init__(self, max_size=0, mode='w+b', bufsize=-1,
  425. suffix="", prefix=template, dir=None):
  426. self._file = _StringIO()
  427. self._max_size = max_size
  428. self._rolled = False
  429. self._TemporaryFileArgs = (mode, bufsize, suffix, prefix, dir)
  430. def _check(self, file):
  431. if self._rolled: return
  432. max_size = self._max_size
  433. if max_size and file.tell() > max_size:
  434. self.rollover()
  435. def rollover(self):
  436. if self._rolled: return
  437. file = self._file
  438. newfile = self._file = TemporaryFile(*self._TemporaryFileArgs)
  439. del self._TemporaryFileArgs
  440. newfile.write(file.getvalue())
  441. newfile.seek(file.tell(), 0)
  442. self._rolled = True
  443. # The method caching trick from NamedTemporaryFile
  444. # won't work here, because _file may change from a
  445. # _StringIO instance to a real file. So we list
  446. # all the methods directly.
  447. # Context management protocol
  448. def __enter__(self):
  449. if self._file.closed:
  450. raise ValueError("Cannot enter context with closed file")
  451. return self
  452. def __exit__(self, exc, value, tb):
  453. self._file.close()
  454. # file protocol
  455. def __iter__(self):
  456. return self._file.__iter__()
  457. def close(self):
  458. self._file.close()
  459. @property
  460. def closed(self):
  461. return self._file.closed
  462. def fileno(self):
  463. self.rollover()
  464. return self._file.fileno()
  465. def flush(self):
  466. self._file.flush()
  467. def isatty(self):
  468. return self._file.isatty()
  469. @property
  470. def mode(self):
  471. try:
  472. return self._file.mode
  473. except AttributeError:
  474. return self._TemporaryFileArgs[0]
  475. @property
  476. def name(self):
  477. try:
  478. return self._file.name
  479. except AttributeError:
  480. return None
  481. def next(self):
  482. return self._file.next
  483. def read(self, *args):
  484. return self._file.read(*args)
  485. def readline(self, *args):
  486. return self._file.readline(*args)
  487. def readlines(self, *args):
  488. return self._file.readlines(*args)
  489. def seek(self, *args):
  490. self._file.seek(*args)
  491. @property
  492. def softspace(self):
  493. return self._file.softspace
  494. def tell(self):
  495. return self._file.tell()
  496. def truncate(self):
  497. self._file.truncate()
  498. def write(self, s):
  499. file = self._file
  500. rv = file.write(s)
  501. self._check(file)
  502. return rv
  503. def writelines(self, iterable):
  504. file = self._file
  505. rv = file.writelines(iterable)
  506. self._check(file)
  507. return rv
  508. def xreadlines(self, *args):
  509. if hasattr(self._file, 'xreadlines'): # real file
  510. return iter(self._file)
  511. else: # StringIO()
  512. return iter(self._file.readlines(*args))