Source code for nxpy.core.temp_file

# nxpy.core package ----------------------------------------------------

# Copyright Nicola Musatti 2010 - 2012
# Use, modification, and distribution are subject to the Boost Software
# License, Version 1.0. (See accompanying file LICENSE.txt or copy at
# http://www.boost.org/LICENSE_1_0.txt)

# See http://sourceforge.net/nxpy for library home page. ---------------------

r"""
Temporary files and directories.

"""

import logging
import os
import shutil
import stat
import tempfile

import nxpy.core.file_object


class TempFile(nxpy.core.file_object.WritableFileObject):
    r"""A temporary file that implements the context manager protocol.
    
    The file is removed when the context is exited from.
    
    """
    def __init__(self, *args, **kwargs):
        file = tempfile.NamedTemporaryFile(*args, delete=False, **kwargs)
        super(TempFile, self).__init__(file)
    
    @property
    def name(self):
        return self._file.name

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self._file.close()
        os.remove(self._file.name)
        return False


class TempDir(object):
    r"""A temporary directory that implements the context manager protocol.
    
    The directory is removed when the context is exited from.
    
    """
    def __init__(self, *args, **kwargs):
        self.dir = tempfile.mkdtemp(*args, **kwargs)
        mode = os.stat(self.dir).st_mode
        os.chmod(self.dir, mode | stat.S_IWRITE)

    @property
    def name(self):
        return self.dir

    def _on_error(self, func, path, ex_info):
        try:
            mode = os.stat(path).st_mode
            if not stat.S_ISLNK(mode):
                os.chmod(path, mode | stat.S_IWRITE)
                if stat.S_ISDIR(mode):
                    os.rmdir(path)
                else:
                    os.remove(path)
            return
        except:
            pass
        logging.error(path + ": not removed")
        
    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        shutil.rmtree(self.dir, ignore_errors=False, onerror=self._on_error)
        return False