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
17import sys
18
19
20class appended_to_syspath(object):
21    """A context for appending a directory to sys.path for a second."""
22
23    def __init__(self, dir):
24        self.dir = dir
25
26    def __enter__(self):
27        if self.dir not in sys.path:
28            sys.path.append(self.dir)
29            self.added = True
30        else:
31            self.added = False
32
33    def __exit__(self, type, value, traceback):
34        if self.added:
35            try:
36                sys.path.remove(self.dir)
37            except ValueError:
38                pass
39        # Returning False causes any exceptions to be re-raised.
40        return False
41
42class prepended_to_syspath(object):
43    """A context for prepending a directory to sys.path for a second."""
44
45    def __init__(self, dir):
46        self.dir = dir
47
48    def __enter__(self):
49        if self.dir not in sys.path:
50            sys.path.insert(0,self.dir)
51            self.added = True
52        else:
53            self.added = False
54
55    def __exit__(self, type, value, traceback):
56        if self.added:
57            try:
58                sys.path.remove(self.dir)
59            except ValueError:
60                pass
61        # Returning False causes any exceptions to be re-raised.
62        return False
63