home *** CD-ROM | disk | FTP | other *** search
/ PC World 2001 April / PCWorld_2001-04_cd.bin / Software / TemaCD / webclean / !!!python!!! / BeOpen-Python-2.0.exe / TEMPFILE.PY < prev    next >
Encoding:
Python Source  |  2000-09-28  |  4.5 KB  |  155 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.  
  8. import os
  9.  
  10.  
  11. # Parameters that the caller may set to override the defaults
  12.  
  13. tempdir = None
  14. template = None
  15.  
  16.  
  17. def gettempdir():
  18.     """Function to calculate the directory to use."""
  19.     global tempdir
  20.     if tempdir is not None:
  21.         return tempdir
  22.     try:
  23.         pwd = os.getcwd()
  24.     except (AttributeError, os.error):
  25.         pwd = os.curdir
  26.     attempdirs = ['/var/tmp', '/usr/tmp', '/tmp', pwd]
  27.     if os.name == 'nt':
  28.         attempdirs.insert(0, 'C:\\TEMP')
  29.         attempdirs.insert(0, '\\TEMP')
  30.     elif os.name == 'mac':
  31.         import macfs, MACFS
  32.         try:
  33.              refnum, dirid = macfs.FindFolder(MACFS.kOnSystemDisk,
  34.                                               MACFS.kTemporaryFolderType, 1)
  35.              dirname = macfs.FSSpec((refnum, dirid, '')).as_pathname()
  36.              attempdirs.insert(0, dirname)
  37.         except macfs.error:
  38.             pass
  39.     for envname in 'TMPDIR', 'TEMP', 'TMP':
  40.         if os.environ.has_key(envname):
  41.             attempdirs.insert(0, os.environ[envname])
  42.     testfile = gettempprefix() + 'test'
  43.     for dir in attempdirs:
  44.         try:
  45.            filename = os.path.join(dir, testfile)
  46.            if os.name == 'posix':
  47.                try:
  48.                    fd = os.open(filename, os.O_RDWR|os.O_CREAT|os.O_EXCL, 0700)
  49.                except OSError:
  50.                    pass
  51.                else:
  52.                    fp = os.fdopen(fd, 'w')
  53.                    fp.write('blat')
  54.                    fp.close()
  55.                    os.unlink(filename)
  56.                    del fp, fd
  57.                    tempdir = dir
  58.                    break
  59.            else:
  60.                fp = open(filename, 'w')
  61.                fp.write('blat')
  62.                fp.close()
  63.                os.unlink(filename)
  64.                tempdir = dir
  65.                break
  66.         except IOError:
  67.             pass
  68.     if tempdir is None:
  69.         msg = "Can't find a usable temporary directory amongst " + `attempdirs`
  70.         raise IOError, msg
  71.     return tempdir
  72.  
  73.  
  74. _pid = None
  75.  
  76. def gettempprefix():
  77.     """Function to calculate a prefix of the filename to use."""
  78.     global template, _pid
  79.     if os.name == 'posix' and _pid and _pid != os.getpid():
  80.         # Our pid changed; we must have forked -- zap the template
  81.         template = None
  82.     if template is None:
  83.         if os.name == 'posix':
  84.             _pid = os.getpid()
  85.             template = '@' + `_pid` + '.'
  86.         elif os.name == 'nt':
  87.             template = '~' + `os.getpid()` + '-'
  88.         elif os.name == 'mac':
  89.             template = 'Python-Tmp-'
  90.         else:
  91.             template = 'tmp' # XXX might choose a better one
  92.     return template
  93.  
  94.  
  95. # Counter for generating unique names
  96.  
  97. counter = 0
  98.  
  99.  
  100. def mktemp(suffix=""):
  101.     """User-callable function to return a unique temporary file name."""
  102.     global counter
  103.     dir = gettempdir()
  104.     pre = gettempprefix()
  105.     while 1:
  106.         counter = counter + 1
  107.         file = os.path.join(dir, pre + `counter` + suffix)
  108.         if not os.path.exists(file):
  109.             return file
  110.  
  111.  
  112. class TemporaryFileWrapper:
  113.     """Temporary file wrapper
  114.  
  115.     This class provides a wrapper around files opened for temporary use.
  116.     In particular, it seeks to automatically remove the file when it is
  117.     no longer needed.
  118.     """
  119.     def __init__(self, file, path):
  120.         self.file = file
  121.         self.path = path
  122.  
  123.     def close(self):
  124.         self.file.close()
  125.         os.unlink(self.path)
  126.  
  127.     def __del__(self):
  128.         try: self.close()
  129.         except: pass
  130.  
  131.     def __getattr__(self, name):
  132.         file = self.__dict__['file']
  133.         a = getattr(file, name)
  134.         if type(a) != type(0):
  135.             setattr(self, name, a)
  136.         return a
  137.  
  138.  
  139. def TemporaryFile(mode='w+b', bufsize=-1, suffix=""):
  140.     """Create and return a temporary file (opened read-write by default)."""
  141.     name = mktemp(suffix)
  142.     if os.name == 'posix':
  143.         # Unix -- be very careful
  144.         fd = os.open(name, os.O_RDWR|os.O_CREAT|os.O_EXCL, 0700)
  145.         try:
  146.             os.unlink(name)
  147.             return os.fdopen(fd, mode, bufsize)
  148.         except:
  149.             os.close(fd)
  150.             raise
  151.     else:
  152.         # Non-unix -- can't unlink file that's still open, use wrapper
  153.         file = open(name, mode, bufsize)
  154.         return TemporaryFileWrapper(file, name)
  155.