Bugfix: Change stdoutput catch

hyperspy/tests/drawing/test_horizontal_plotting.py

@@ -19,57 +19,52 @@
 import hyperspy.api as hs
 import numpy as np
 import matplotlib
-from contextlib import redirect_stdout
-import io
+
 ipympl = pytest.importorskip("ipympl")
 ipywidgets = pytest.importorskip("ipywidgets")
 from ipywidgets import Output
+from IPython.utils.capture import capture_output
 
 class TestIPYMPL:
     def test_horizontal(self):
         matplotlib.use("module://ipympl.backend_nbagg")
-        f = io.StringIO()
-        with redirect_stdout(f):
+        with capture_output() as c:
             s = hs.signals.Signal2D(np.random.random((4, 4, 2, 2)))
             s.plot(plot_style="horizontal")
-
         assert("HBox(children=(Output(layout=Layout(margin='auto 0px auto 0px')),"
-               " Output(layout=Layout(margin='auto 0px auto 0px')" in f.getvalue())
+               " Output(layout=Layout(margin='auto 0px auto 0px')" in c.stdout)
 
     def test_vertical(self):
         matplotlib.use("module://ipympl.backend_nbagg")
-        f = io.StringIO()
-        with redirect_stdout(f):
+        with capture_output() as c:
             s = hs.signals.Signal2D(np.random.random((4, 4, 2, 2)))
             s.plot(plot_style="vertical")
 
         assert("VBox(children=(Output(layout=Layout(margin='0px auto 0px auto')),"
-               " Output(layout=Layout(margin='0px auto 0px auto')" in f.getvalue())
+               " Output(layout=Layout(margin='0px auto 0px auto')" in c.stdout)
 
     def test_one_output(self):
         matplotlib.use("module://ipympl.backend_nbagg")
-        f = io.StringIO()
-        with redirect_stdout(f):
+        with capture_output() as c:
             o = Output()
             s = hs.signals.Signal2D(np.random.random((4, 4, 2, 2)))
             s.plot(widget=o)
-        assert("Output()" in f.getvalue())
+        assert("Output()" in c.stdout)
 
-        with redirect_stdout(f):
+        with capture_output() as c:
             o = Output()
             s = hs.signals.Signal2D(np.random.random((4, 4, 2, 2)))
             s.plot(navigator_kwds={"widget":o})
-        assert("Output()" in f.getvalue())
+        assert("Output()" in c.stdout)
 
 
     def test_two_outputs(self):
         matplotlib.use("module://ipympl.backend_nbagg")
-        f = io.StringIO()
-        with redirect_stdout(f):
+        with capture_output() as c:
             o = Output()
             o2 = Output()
             s = hs.signals.Signal2D(np.random.random((4, 4, 2, 2)))
             s.plot(widget=o, navigator_kwds={"widget":o2})
-        assert("Output" not in f.getvalue())
+        assert("Output" not in c.stdout)