display_trap.py 2.0 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970
  1. # encoding: utf-8
  2. """
  3. A context manager for handling sys.displayhook.
  4. Authors:
  5. * Robert Kern
  6. * Brian Granger
  7. """
  8. #-----------------------------------------------------------------------------
  9. # Copyright (C) 2008-2011 The IPython Development Team
  10. #
  11. # Distributed under the terms of the BSD License. The full license is in
  12. # the file COPYING, distributed as part of this software.
  13. #-----------------------------------------------------------------------------
  14. #-----------------------------------------------------------------------------
  15. # Imports
  16. #-----------------------------------------------------------------------------
  17. import sys
  18. from traitlets.config.configurable import Configurable
  19. from traitlets import Any
  20. #-----------------------------------------------------------------------------
  21. # Classes and functions
  22. #-----------------------------------------------------------------------------
  23. class DisplayTrap(Configurable):
  24. """Object to manage sys.displayhook.
  25. This came from IPython.core.kernel.display_hook, but is simplified
  26. (no callbacks or formatters) until more of the core is refactored.
  27. """
  28. hook = Any()
  29. def __init__(self, hook=None):
  30. super(DisplayTrap, self).__init__(hook=hook, config=None)
  31. self.old_hook = None
  32. # We define this to track if a single BuiltinTrap is nested.
  33. # Only turn off the trap when the outermost call to __exit__ is made.
  34. self._nested_level = 0
  35. def __enter__(self):
  36. if self._nested_level == 0:
  37. self.set()
  38. self._nested_level += 1
  39. return self
  40. def __exit__(self, type, value, traceback):
  41. if self._nested_level == 1:
  42. self.unset()
  43. self._nested_level -= 1
  44. # Returning False will cause exceptions to propagate
  45. return False
  46. def set(self):
  47. """Set the hook."""
  48. if sys.displayhook is not self.hook:
  49. self.old_hook = sys.displayhook
  50. sys.displayhook = self.hook
  51. def unset(self):
  52. """Unset the hook."""
  53. sys.displayhook = self.old_hook