IPython/utils/syspathcontext.py

# encoding: utf-8
"""
Context managers for adding things to sys.path temporarily.

Authors:

* Brian Granger
"""

#-----------------------------------------------------------------------------
#  Copyright (C) 2008-2011  The IPython Development Team
#
#  Distributed under the terms of the BSD License.  The full license is in
#  the file COPYING, distributed as part of this software.
#-----------------------------------------------------------------------------

#-----------------------------------------------------------------------------
# Imports
#-----------------------------------------------------------------------------

import sys

from IPython.utils.py3compat import cast_bytes_py2

#-----------------------------------------------------------------------------
# Code
#-----------------------------------------------------------------------------

class appended_to_syspath(object):
    """A context for appending a directory to sys.path for a second."""

    def __init__(self, dir):
        self.dir = cast_bytes_py2(dir, sys.getdefaultencoding())

    def __enter__(self):
        if self.dir not in sys.path:
            sys.path.append(self.dir)
            self.added = True
        else:
            self.added = False

    def __exit__(self, type, value, traceback):
        if self.added:
            try:
                sys.path.remove(self.dir)
            except ValueError:
                pass
        # Returning False causes any exceptions to be re-raised.
        return False

class prepended_to_syspath(object):
    """A context for prepending a directory to sys.path for a second."""

    def __init__(self, dir):
        self.dir = cast_bytes_py2(dir, sys.getdefaultencoding())

    def __enter__(self):
        if self.dir not in sys.path:
            sys.path.insert(0,self.dir)
            self.added = True
        else:
            self.added = False

    def __exit__(self, type, value, traceback):
        if self.added:
            try:
                sys.path.remove(self.dir)
            except ValueError:
                pass
        # Returning False causes any exceptions to be re-raised.
        return False