| 18 | |
| 19 | |
| 20 | class BuiltinTrap(Configurable): |
| 21 | |
| 22 | shell = Instance('IPython.core.interactiveshell.InteractiveShellABC', |
| 23 | allow_none=True) |
| 24 | |
| 25 | def __init__(self, shell=None): |
| 26 | super(BuiltinTrap, self).__init__(shell=shell, config=None) |
| 27 | self._orig_builtins = {} |
| 28 | # We define this to track if a single BuiltinTrap is nested. |
| 29 | # Only turn off the trap when the outermost call to __exit__ is made. |
| 30 | self._nested_level = 0 |
| 31 | self.shell = shell |
| 32 | # builtins we always add - if set to HideBuiltin, they will just |
| 33 | # be removed instead of being replaced by something else |
| 34 | self.auto_builtins = {'exit': HideBuiltin, |
| 35 | 'quit': HideBuiltin, |
| 36 | 'get_ipython': self.shell.get_ipython, |
| 37 | } |
| 38 | |
| 39 | def __enter__(self): |
| 40 | if self._nested_level == 0: |
| 41 | self.activate() |
| 42 | self._nested_level += 1 |
| 43 | # I return self, so callers can use add_builtin in a with clause. |
| 44 | return self |
| 45 | |
| 46 | def __exit__(self, type, value, traceback): |
| 47 | if self._nested_level == 1: |
| 48 | self.deactivate() |
| 49 | self._nested_level -= 1 |
| 50 | # Returning False will cause exceptions to propagate |
| 51 | return False |
| 52 | |
| 53 | def add_builtin(self, key, value): |
| 54 | """Add a builtin and save the original.""" |
| 55 | bdict = builtin_mod.__dict__ |
| 56 | orig = bdict.get(key, BuiltinUndefined) |
| 57 | if value is HideBuiltin: |
| 58 | if orig is not BuiltinUndefined: #same as 'key in bdict' |
| 59 | self._orig_builtins[key] = orig |
| 60 | del bdict[key] |
| 61 | else: |
| 62 | self._orig_builtins[key] = orig |
| 63 | bdict[key] = value |
| 64 | |
| 65 | def remove_builtin(self, key, orig): |
| 66 | """Remove an added builtin and re-set the original.""" |
| 67 | if orig is BuiltinUndefined: |
| 68 | del builtin_mod.__dict__[key] |
| 69 | else: |
| 70 | builtin_mod.__dict__[key] = orig |
| 71 | |
| 72 | def activate(self): |
| 73 | """Store ipython references in the __builtin__ namespace.""" |
| 74 | |
| 75 | add_builtin = self.add_builtin |
| 76 | for name, func in self.auto_builtins.items(): |
| 77 | add_builtin(name, func) |