1# encoding: utf-8
2"""
3Context managers for adding things to sys.path temporarily.
4
5Authors:
6
7* Brian Granger
8"""
9
10#-----------------------------------------------------------------------------
11#  Copyright (C) 2008-2011  The IPython Development Team
12#
13#  Distributed under the terms of the BSD License.  The full license is in
14#  the file COPYING, distributed as part of this software.
15#-----------------------------------------------------------------------------
16
17#-----------------------------------------------------------------------------
18# Imports
19#-----------------------------------------------------------------------------
20
21import sys
22
23from IPython.utils.py3compat import cast_bytes_py2
24
25#-----------------------------------------------------------------------------
26# Code
27#-----------------------------------------------------------------------------
28
29class appended_to_syspath(object):
30    """A context for appending a directory to sys.path for a second."""
31
32    def __init__(self, dir):
33        self.dir = cast_bytes_py2(dir, sys.getdefaultencoding())
34
35    def __enter__(self):
36        if self.dir not in sys.path:
37            sys.path.append(self.dir)
38            self.added = True
39        else:
40            self.added = False
41
42    def __exit__(self, type, value, traceback):
43        if self.added:
44            try:
45                sys.path.remove(self.dir)
46            except ValueError:
47                pass
48        # Returning False causes any exceptions to be re-raised.
49        return False
50
51class prepended_to_syspath(object):
52    """A context for prepending a directory to sys.path for a second."""
53
54    def __init__(self, dir):
55        self.dir = cast_bytes_py2(dir, sys.getdefaultencoding())
56
57    def __enter__(self):
58        if self.dir not in sys.path:
59            sys.path.insert(0,self.dir)
60            self.added = True
61        else:
62            self.added = False
63
64    def __exit__(self, type, value, traceback):
65        if self.added:
66            try:
67                sys.path.remove(self.dir)
68            except ValueError:
69                pass
70        # Returning False causes any exceptions to be re-raised.
71        return False
72