context manager for capturing stdout/err
| 122 | |
| 123 | |
| 124 | class capture_output(object): |
| 125 | """context manager for capturing stdout/err""" |
| 126 | stdout = True |
| 127 | stderr = True |
| 128 | display = True |
| 129 | |
| 130 | def __init__(self, stdout=True, stderr=True, display=True): |
| 131 | self.stdout = stdout |
| 132 | self.stderr = stderr |
| 133 | self.display = display |
| 134 | self.shell = None |
| 135 | |
| 136 | def __enter__(self): |
| 137 | from IPython.core.getipython import get_ipython |
| 138 | from IPython.core.displaypub import CapturingDisplayPublisher |
| 139 | from IPython.core.displayhook import CapturingDisplayHook |
| 140 | |
| 141 | self.sys_stdout = sys.stdout |
| 142 | self.sys_stderr = sys.stderr |
| 143 | |
| 144 | if self.display: |
| 145 | self.shell = get_ipython() |
| 146 | if self.shell is None: |
| 147 | self.save_display_pub = None |
| 148 | self.display = False |
| 149 | |
| 150 | stdout = stderr = outputs = None |
| 151 | if self.stdout: |
| 152 | stdout = sys.stdout = StringIO() |
| 153 | if self.stderr: |
| 154 | stderr = sys.stderr = StringIO() |
| 155 | if self.display: |
| 156 | self.save_display_pub = self.shell.display_pub |
| 157 | self.shell.display_pub = CapturingDisplayPublisher() |
| 158 | outputs = self.shell.display_pub.outputs |
| 159 | self.save_display_hook = sys.displayhook |
| 160 | sys.displayhook = CapturingDisplayHook(shell=self.shell, |
| 161 | outputs=outputs) |
| 162 | |
| 163 | return CapturedIO(stdout, stderr, outputs) |
| 164 | |
| 165 | def __exit__(self, exc_type, exc_value, traceback): |
| 166 | sys.stdout = self.sys_stdout |
| 167 | sys.stderr = self.sys_stderr |
| 168 | if self.display and self.shell: |
| 169 | self.shell.display_pub = self.save_display_pub |
| 170 | sys.displayhook = self.save_display_hook |
no outgoing calls