home *** CD-ROM | disk | FTP | other *** search
/ PC Extra 07 & 08 / pca1507.iso / Software / psp8 / Data1.cab / tempfile.py < prev    next >
Encoding:
Python Source  |  2003-04-22  |  6.9 KB  |  219 lines
  1. """Temporary files and filenames."""
  2.  
  3. # XXX This tries to be not UNIX specific, but I don't know beans about
  4. # how to choose a temp directory or filename on MS-DOS or other
  5. # systems so it may have to be changed...
  6.  
  7. import os
  8.  
  9. __all__ = ["mktemp", "TemporaryFile", "tempdir", "gettempprefix"]
  10.  
  11. # Parameters that the caller may set to override the defaults
  12. tempdir = None
  13. template = None
  14.  
  15. def gettempdir():
  16.     """Function to calculate the directory to use."""
  17.     global tempdir
  18.     if tempdir is not None:
  19.         return tempdir
  20.     try:
  21.         pwd = os.getcwd()
  22.     except (AttributeError, os.error):
  23.         pwd = os.curdir
  24.     attempdirs = ['/tmp', '/var/tmp', '/usr/tmp', pwd]
  25.     if os.name == 'nt':
  26.         attempdirs.insert(0, 'C:\\TEMP')
  27.         attempdirs.insert(0, '\\TEMP')
  28.     elif os.name == 'mac':
  29.         import macfs, MACFS
  30.         try:
  31.             refnum, dirid = macfs.FindFolder(MACFS.kOnSystemDisk,
  32.                                              MACFS.kTemporaryFolderType, 1)
  33.             dirname = macfs.FSSpec((refnum, dirid, '')).as_pathname()
  34.             attempdirs.insert(0, dirname)
  35.         except macfs.error:
  36.             pass
  37.     elif os.name == 'riscos':
  38.         scrapdir = os.getenv('Wimp$ScrapDir')
  39.         if scrapdir:
  40.             attempdirs.insert(0, scrapdir)
  41.     for envname in 'TMPDIR', 'TEMP', 'TMP':
  42.         if os.environ.has_key(envname):
  43.             attempdirs.insert(0, os.environ[envname])
  44.     testfile = gettempprefix() + 'test'
  45.     for dir in attempdirs:
  46.         try:
  47.             filename = os.path.join(dir, testfile)
  48.             if os.name == 'posix':
  49.                 try:
  50.                     fd = os.open(filename,
  51.                                  os.O_RDWR | os.O_CREAT | os.O_EXCL, 0700)
  52.                 except OSError:
  53.                     pass
  54.                 else:
  55.                     fp = os.fdopen(fd, 'w')
  56.                     fp.write('blat')
  57.                     fp.close()
  58.                     os.unlink(filename)
  59.                     del fp, fd
  60.                     tempdir = dir
  61.                     break
  62.             else:
  63.                 fp = open(filename, 'w')
  64.                 fp.write('blat')
  65.                 fp.close()
  66.                 os.unlink(filename)
  67.                 tempdir = dir
  68.                 break
  69.         except IOError:
  70.             pass
  71.     if tempdir is None:
  72.         msg = "Can't find a usable temporary directory amongst " + `attempdirs`
  73.         raise IOError, msg
  74.     return tempdir
  75.  
  76.  
  77. # template caches the result of gettempprefix, for speed, when possible.
  78. # XXX unclear why this isn't "_template"; left it "template" for backward
  79. # compatibility.
  80. if os.name == "posix":
  81.     # We don't try to cache the template on posix:  the pid may change on us
  82.     # between calls due to a fork, and on Linux the pid changes even for
  83.     # another thread in the same process.  Since any attempt to keep the
  84.     # cache in synch would have to call os.getpid() anyway in order to make
  85.     # sure the pid hasn't changed between calls, a cache wouldn't save any
  86.     # time.  In addition, a cache is difficult to keep correct with the pid
  87.     # changing willy-nilly, and earlier attempts proved buggy (races).
  88.     template = None
  89.  
  90. # Else the pid never changes, so gettempprefix always returns the same
  91. # string.
  92. elif os.name == "nt":
  93.     template = '~' + `os.getpid()` + '-'
  94. elif os.name in ('mac', 'riscos'):
  95.     template = 'Python-Tmp-'
  96. else:
  97.     template = 'tmp' # XXX might choose a better one
  98.  
  99. def gettempprefix():
  100.     """Function to calculate a prefix of the filename to use.
  101.  
  102.     This incorporates the current process id on systems that support such a
  103.     notion, so that concurrent processes don't generate the same prefix.
  104.     """
  105.  
  106.     global template
  107.     if template is None:
  108.         return '@' + `os.getpid()` + '.'
  109.     else:
  110.         return template
  111.  
  112.  
  113. def mktemp(suffix=""):
  114.     """User-callable function to return a unique temporary file name."""
  115.     dir = gettempdir()
  116.     pre = gettempprefix()
  117.     while 1:
  118.         i = _counter.get_next()
  119.         file = os.path.join(dir, pre + str(i) + suffix)
  120.         if not os.path.exists(file):
  121.             return file
  122.  
  123.  
  124. class TemporaryFileWrapper:
  125.     """Temporary file wrapper
  126.  
  127.     This class provides a wrapper around files opened for temporary use.
  128.     In particular, it seeks to automatically remove the file when it is
  129.     no longer needed.
  130.     """
  131.  
  132.     # Cache the unlinker so we don't get spurious errors at shutdown
  133.     # when the module-level "os" is None'd out.  Note that this must
  134.     # be referenced as self.unlink, because the name TemporaryFileWrapper
  135.     # may also get None'd out before __del__ is called.
  136.     unlink = os.unlink
  137.  
  138.     def __init__(self, file, path):
  139.         self.file = file
  140.         self.path = path
  141.         self.close_called = 0
  142.  
  143.     def close(self):
  144.         if not self.close_called:
  145.             self.close_called = 1
  146.             self.file.close()
  147.             self.unlink(self.path)
  148.  
  149.     def __del__(self):
  150.         self.close()
  151.  
  152.     def __getattr__(self, name):
  153.         file = self.__dict__['file']
  154.         a = getattr(file, name)
  155.         if type(a) != type(0):
  156.             setattr(self, name, a)
  157.         return a
  158.  
  159.  
  160. def TemporaryFile(mode='w+b', bufsize=-1, suffix=""):
  161.     """Create and return a temporary file (opened read-write by default)."""
  162.     name = mktemp(suffix)
  163.     if os.name == 'posix':
  164.         # Unix -- be very careful
  165.         fd = os.open(name, os.O_RDWR|os.O_CREAT|os.O_EXCL, 0700)
  166.         try:
  167.             os.unlink(name)
  168.             return os.fdopen(fd, mode, bufsize)
  169.         except:
  170.             os.close(fd)
  171.             raise
  172.     else:
  173.         # Non-unix -- can't unlink file that's still open, use wrapper
  174.         file = open(name, mode, bufsize)
  175.         return TemporaryFileWrapper(file, name)
  176.  
  177. # In order to generate unique names, mktemp() uses _counter.get_next().
  178. # This returns a unique integer on each call, in a threadsafe way (i.e.,
  179. # multiple threads will never see the same integer).  The integer will
  180. # usually be a Python int, but if _counter.get_next() is called often
  181. # enough, it will become a Python long.
  182. # Note that the only name that survives this next block of code
  183. # is "_counter".
  184.  
  185. class _ThreadSafeCounter:
  186.     def __init__(self, mutex, initialvalue=0):
  187.         self.mutex = mutex
  188.         self.i = initialvalue
  189.  
  190.     def get_next(self):
  191.         self.mutex.acquire()
  192.         result = self.i
  193.         try:
  194.             newi = result + 1
  195.         except OverflowError:
  196.             newi = long(result) + 1
  197.         self.i = newi
  198.         self.mutex.release()
  199.         return result
  200.  
  201. try:
  202.     import thread
  203.  
  204. except ImportError:
  205.     class _DummyMutex:
  206.         def acquire(self):
  207.             pass
  208.  
  209.         release = acquire
  210.  
  211.     _counter = _ThreadSafeCounter(_DummyMutex())
  212.     del _DummyMutex
  213.  
  214. else:
  215.     _counter = _ThreadSafeCounter(thread.allocate_lock())
  216.     del thread
  217.  
  218. del _ThreadSafeCounter
  219.