diff --git a/core/ipyflow/utils/ipython_utils.py b/core/ipyflow/utils/ipython_utils.py index 8a6c1856..8a8d9e9e 100644 --- a/core/ipyflow/utils/ipython_utils.py +++ b/core/ipyflow/utils/ipython_utils.py @@ -150,6 +150,11 @@ def set_parent(self, *args, **kwargs) -> None: self.pub2.set_parent(*args, **kwargs) +class TeeCompatibleCapturingDisplayPublisher(CapturingDisplayPublisher): + def clear_output(self, wait=False): + self.outputs.clear() + + class CaptureOutputTee: """ Context manager for capturing and replicating stdout/err and rich display publishers. @@ -186,7 +191,7 @@ def __enter__(self) -> CapturedIO: sys.stderr = Tee(sys.stderr, stderr) # type: ignore if self.display and self.shell is not None: self.save_display_pub = self.shell.display_pub - capture_display_pub = CapturingDisplayPublisher() + capture_display_pub = TeeCompatibleCapturingDisplayPublisher() outputs = capture_display_pub.outputs self.shell.display_pub = TeeDisplayPublisher( self.save_display_pub, capture_display_pub