home *** CD-ROM | disk | FTP | other *** search
/ H4CK3R 14 / hacker14.iso / programacao / pythonwin / python.exe / TEMPFILE.PY < prev    next >
Encoding:
Python Source  |  2003-07-22  |  14.1 KB  |  452 lines
  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.     # If PYTHONCASEOK is set on Windows, stinking FCNTL.py gets
  42.     # imported, and we don't get an ImportError then.  Provoke
  43.     # an AttributeError instead in that case.
  44.     _fcntl.fcntl
  45. except (ImportError, AttributeError):
  46.     def _set_cloexec(fd):
  47.         pass
  48. else:
  49.     def _set_cloexec(fd):
  50.         flags = _fcntl.fcntl(fd, _fcntl.F_GETFD, 0)
  51.         if flags >= 0:
  52.             # flags read successfully, modify
  53.             flags |= _fcntl.FD_CLOEXEC
  54.             _fcntl.fcntl(fd, _fcntl.F_SETFD, flags)
  55.  
  56.  
  57. try:
  58.     import thread as _thread
  59. except ImportError:
  60.     import dummy_thread as _thread
  61. _allocate_lock = _thread.allocate_lock
  62.  
  63. _text_openflags = _os.O_RDWR | _os.O_CREAT | _os.O_EXCL
  64. if hasattr(_os, 'O_NOINHERIT'):
  65.     _text_openflags |= _os.O_NOINHERIT
  66. if hasattr(_os, 'O_NOFOLLOW'):
  67.     _text_openflags |= _os.O_NOFOLLOW
  68.  
  69. _bin_openflags = _text_openflags
  70. if hasattr(_os, 'O_BINARY'):
  71.     _bin_openflags |= _os.O_BINARY
  72.  
  73. if hasattr(_os, 'TMP_MAX'):
  74.     TMP_MAX = _os.TMP_MAX
  75. else:
  76.     TMP_MAX = 10000
  77.  
  78. template = "tmp"
  79.  
  80. tempdir = None
  81.  
  82. # Internal routines.
  83.  
  84. _once_lock = _allocate_lock()
  85.  
  86. class _RandomNameSequence:
  87.     """An instance of _RandomNameSequence generates an endless
  88.     sequence of unpredictable strings which can safely be incorporated
  89.     into file names.  Each string is six characters long.  Multiple
  90.     threads can safely use the same instance at the same time.
  91.  
  92.     _RandomNameSequence is an iterator."""
  93.  
  94.     characters = ("abcdefghijklmnopqrstuvwxyz" +
  95.                   "ABCDEFGHIJKLMNOPQRSTUVWXYZ" +
  96.                   "0123456789-_")
  97.  
  98.     def __init__(self):
  99.         self.mutex = _allocate_lock()
  100.         self.rng = _Random()
  101.         self.normcase = _os.path.normcase
  102.  
  103.     def __iter__(self):
  104.         return self
  105.  
  106.     def next(self):
  107.         m = self.mutex
  108.         c = self.characters
  109.         choose = self.rng.choice
  110.  
  111.         m.acquire()
  112.         try:
  113.             letters = [choose(c) for dummy in "123456"]
  114.         finally:
  115.             m.release()
  116.  
  117.         return self.normcase(''.join(letters))
  118.  
  119. def _candidate_tempdir_list():
  120.     """Generate a list of candidate temporary directories which
  121.     _get_default_tempdir will try."""
  122.  
  123.     dirlist = []
  124.  
  125.     # First, try the environment.
  126.     for envname in 'TMPDIR', 'TEMP', 'TMP':
  127.         dirname = _os.getenv(envname)
  128.         if dirname: dirlist.append(dirname)
  129.  
  130.     # Failing that, try OS-specific locations.
  131.     if _os.name == 'mac':
  132.         try:
  133.             fsr = _Folder.FSFindFolder(_Folders.kOnSystemDisk,
  134.                                               _Folders.kTemporaryFolderType, 1)
  135.             dirname = fsr.as_pathname()
  136.             dirlist.append(dirname)
  137.         except _Folder.error:
  138.             pass
  139.     elif _os.name == 'riscos':
  140.         dirname = _os.getenv('Wimp$ScrapDir')
  141.         if dirname: dirlist.append(dirname)
  142.     elif _os.name == 'nt':
  143.         dirlist.extend([ r'c:\temp', r'c:\tmp', r'\temp', r'\tmp' ])
  144.     else:
  145.         dirlist.extend([ '/tmp', '/var/tmp', '/usr/tmp' ])
  146.  
  147.     # As a last resort, the current directory.
  148.     try:
  149.         dirlist.append(_os.getcwd())
  150.     except (AttributeError, _os.error):
  151.         dirlist.append(_os.curdir)
  152.  
  153.     return dirlist
  154.  
  155. def _get_default_tempdir():
  156.     """Calculate the default directory to use for temporary files.
  157.     This routine should be called exactly once.
  158.  
  159.     We determine whether or not a candidate temp dir is usable by
  160.     trying to create and write to a file in that directory.  If this
  161.     is successful, the test file is deleted.  To prevent denial of
  162.     service, the name of the test file must be randomized."""
  163.  
  164.     namer = _RandomNameSequence()
  165.     dirlist = _candidate_tempdir_list()
  166.     flags = _text_openflags
  167.  
  168.     for dir in dirlist:
  169.         if dir != _os.curdir:
  170.             dir = _os.path.normcase(_os.path.abspath(dir))
  171.         # Try only a few names per directory.
  172.         for seq in xrange(100):
  173.             name = namer.next()
  174.             filename = _os.path.join(dir, name)
  175.             try:
  176.                 fd = _os.open(filename, flags, 0600)
  177.                 fp = _os.fdopen(fd, 'w')
  178.                 fp.write('blat')
  179.                 fp.close()
  180.                 _os.unlink(filename)
  181.                 del fp, fd
  182.                 return dir
  183.             except (OSError, IOError), e:
  184.                 if e[0] != _errno.EEXIST:
  185.                     break # no point trying more names in this directory
  186.                 pass
  187.     raise IOError, (_errno.ENOENT,
  188.                     ("No usable temporary directory found in %s" % dirlist))
  189.  
  190. _name_sequence = None
  191.  
  192. def _get_candidate_names():
  193.     """Common setup sequence for all user-callable interfaces."""
  194.  
  195.     global _name_sequence
  196.     if _name_sequence is None:
  197.         _once_lock.acquire()
  198.         try:
  199.             if _name_sequence is None:
  200.                 _name_sequence = _RandomNameSequence()
  201.         finally:
  202.             _once_lock.release()
  203.     return _name_sequence
  204.  
  205.  
  206. def _mkstemp_inner(dir, pre, suf, flags):
  207.     """Code common to mkstemp, TemporaryFile, and NamedTemporaryFile."""
  208.  
  209.     names = _get_candidate_names()
  210.  
  211.     for seq in xrange(TMP_MAX):
  212.         name = names.next()
  213.         file = _os.path.join(dir, pre + name + suf)
  214.         try:
  215.             fd = _os.open(file, flags, 0600)
  216.             _set_cloexec(fd)
  217.             return (fd, file)
  218.         except OSError, e:
  219.             if e.errno == _errno.EEXIST:
  220.                 continue # try again
  221.             raise
  222.  
  223.     raise IOError, (_errno.EEXIST, "No usable temporary file name found")
  224.  
  225.  
  226. # User visible interfaces.
  227.  
  228. def gettempprefix():
  229.     """Accessor for tempdir.template."""
  230.     return template
  231.  
  232. tempdir = None
  233.  
  234. def gettempdir():
  235.     """Accessor for tempdir.tempdir."""
  236.     global tempdir
  237.     if tempdir is None:
  238.         _once_lock.acquire()
  239.         try:
  240.             if tempdir is None:
  241.                 tempdir = _get_default_tempdir()
  242.         finally:
  243.             _once_lock.release()
  244.     return tempdir
  245.  
  246. def mkstemp(suffix="", prefix=template, dir=None, text=False):
  247.     """mkstemp([suffix, [prefix, [dir, [text]]]])
  248.     User-callable function to create and return a unique temporary
  249.     file.  The return value is a pair (fd, name) where fd is the
  250.     file descriptor returned by os.open, and name is the filename.
  251.  
  252.     If 'suffix' is specified, the file name will end with that suffix,
  253.     otherwise there will be no suffix.
  254.  
  255.     If 'prefix' is specified, the file name will begin with that prefix,
  256.     otherwise a default prefix is used.
  257.  
  258.     If 'dir' is specified, the file will be created in that directory,
  259.     otherwise a default directory is used.
  260.  
  261.     If 'text' is specified and true, the file is opened in text
  262.     mode.  Else (the default) the file is opened in binary mode.  On
  263.     some operating systems, this makes no difference.
  264.  
  265.     The file is readable and writable only by the creating user ID.
  266.     If the operating system uses permission bits to indicate whether a
  267.     file is executable, the file is executable by no one. The file
  268.     descriptor is not inherited by children of this process.
  269.  
  270.     Caller is responsible for deleting the file when done with it.
  271.     """
  272.  
  273.     if dir is None:
  274.         dir = gettempdir()
  275.  
  276.     if text:
  277.         flags = _text_openflags
  278.     else:
  279.         flags = _bin_openflags
  280.  
  281.     return _mkstemp_inner(dir, prefix, suffix, flags)
  282.  
  283.  
  284. def mkdtemp(suffix="", prefix=template, dir=None):
  285.     """mkdtemp([suffix, [prefix, [dir]]])
  286.     User-callable function to create and return a unique temporary
  287.     directory.  The return value is the pathname of the directory.
  288.  
  289.     Arguments are as for mkstemp, except that the 'text' argument is
  290.     not accepted.
  291.  
  292.     The directory is readable, writable, and searchable only by the
  293.     creating user.
  294.  
  295.     Caller is responsible for deleting the directory when done with it.
  296.     """
  297.  
  298.     if dir is None:
  299.         dir = gettempdir()
  300.  
  301.     names = _get_candidate_names()
  302.  
  303.     for seq in xrange(TMP_MAX):
  304.         name = names.next()
  305.         file = _os.path.join(dir, prefix + name + suffix)
  306.         try:
  307.             _os.mkdir(file, 0700)
  308.             return file
  309.         except OSError, e:
  310.             if e.errno == _errno.EEXIST:
  311.                 continue # try again
  312.             raise
  313.  
  314.     raise IOError, (_errno.EEXIST, "No usable temporary directory name found")
  315.  
  316. def mktemp(suffix="", prefix=template, dir=None):
  317.     """mktemp([suffix, [prefix, [dir]]])
  318.     User-callable function to return a unique temporary file name.  The
  319.     file is not created.
  320.  
  321.     Arguments are as for mkstemp, except that the 'text' argument is
  322.     not accepted.
  323.  
  324.     This function is unsafe and should not be used.  The file name
  325.     refers to a file that did not exist at some point, but by the time
  326.     you get around to creating it, someone else may have beaten you to
  327.     the punch.
  328.     """
  329.  
  330. ##    from warnings import warn as _warn
  331. ##    _warn("mktemp is a potential security risk to your program",
  332. ##          RuntimeWarning, stacklevel=2)
  333.  
  334.     if dir is None:
  335.         dir = gettempdir()
  336.  
  337.     names = _get_candidate_names()
  338.     for seq in xrange(TMP_MAX):
  339.         name = names.next()
  340.         file = _os.path.join(dir, prefix + name + suffix)
  341.         if not _os.path.exists(file):
  342.             return file
  343.  
  344.     raise IOError, (_errno.EEXIST, "No usable temporary filename found")
  345.  
  346. class _TemporaryFileWrapper:
  347.     """Temporary file wrapper
  348.  
  349.     This class provides a wrapper around files opened for
  350.     temporary use.  In particular, it seeks to automatically
  351.     remove the file when it is no longer needed.
  352.     """
  353.  
  354.     def __init__(self, file, name):
  355.         self.file = file
  356.         self.name = name
  357.         self.close_called = False
  358.  
  359.     def __getattr__(self, name):
  360.         file = self.__dict__['file']
  361.         a = getattr(file, name)
  362.         if type(a) != type(0):
  363.             setattr(self, name, a)
  364.         return a
  365.  
  366.     # NT provides delete-on-close as a primitive, so we don't need
  367.     # the wrapper to do anything special.  We still use it so that
  368.     # file.name is useful (i.e. not "(fdopen)") with NamedTemporaryFile.
  369.     if _os.name != 'nt':
  370.  
  371.         # Cache the unlinker so we don't get spurious errors at
  372.         # shutdown when the module-level "os" is None'd out.  Note
  373.         # that this must be referenced as self.unlink, because the
  374.         # name TemporaryFileWrapper may also get None'd out before
  375.         # __del__ is called.
  376.         unlink = _os.unlink
  377.  
  378.         def close(self):
  379.             if not self.close_called:
  380.                 self.close_called = True
  381.                 self.file.close()
  382.                 self.unlink(self.name)
  383.  
  384.         def __del__(self):
  385.             self.close()
  386.  
  387. def NamedTemporaryFile(mode='w+b', bufsize=-1, suffix="",
  388.                        prefix=template, dir=None):
  389.     """Create and return a temporary file.
  390.     Arguments:
  391.     'prefix', 'suffix', 'dir' -- as for mkstemp.
  392.     'mode' -- the mode argument to os.fdopen (default "w+b").
  393.     'bufsize' -- the buffer size argument to os.fdopen (default -1).
  394.     The file is created as mkstemp() would do it.
  395.  
  396.     Returns a file object; the name of the file is accessible as
  397.     file.name.  The file will be automatically deleted when it is
  398.     closed.
  399.     """
  400.  
  401.     if dir is None:
  402.         dir = gettempdir()
  403.  
  404.     if 'b' in mode:
  405.         flags = _bin_openflags
  406.     else:
  407.         flags = _text_openflags
  408.  
  409.     # Setting O_TEMPORARY in the flags causes the OS to delete
  410.     # the file when it is closed.  This is only supported by Windows.
  411.     if _os.name == 'nt':
  412.         flags |= _os.O_TEMPORARY
  413.  
  414.     (fd, name) = _mkstemp_inner(dir, prefix, suffix, flags)
  415.     file = _os.fdopen(fd, mode, bufsize)
  416.     return _TemporaryFileWrapper(file, name)
  417.  
  418. if _os.name != 'posix' or _os.sys.platform == 'cygwin':
  419.     # On non-POSIX and Cygwin systems, assume that we cannot unlink a file
  420.     # while it is open.
  421.     TemporaryFile = NamedTemporaryFile
  422.  
  423. else:
  424.     def TemporaryFile(mode='w+b', bufsize=-1, suffix="",
  425.                       prefix=template, dir=None):
  426.         """Create and return a temporary file.
  427.         Arguments:
  428.         'prefix', 'suffix', 'directory' -- 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 a file object.  The file has no name, and will cease to
  434.         exist when it is closed.
  435.         """
  436.  
  437.         if dir is None:
  438.             dir = gettempdir()
  439.  
  440.         if 'b' in mode:
  441.             flags = _bin_openflags
  442.         else:
  443.             flags = _text_openflags
  444.  
  445.         (fd, name) = _mkstemp_inner(dir, prefix, suffix, flags)
  446.         try:
  447.             _os.unlink(name)
  448.             return _os.fdopen(fd, mode, bufsize)
  449.         except:
  450.             _os.close(fd)
  451.             raise
  452.