ensure final state is same as initial for ion/ioff contexts

ianhi · ianhi · commit d4780709d98e · 2020-08-11T21:15:52.000-04:00
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -385,6 +385,9 @@ def __exit__(self, exc_type, exc_value, traceback):
         if self.wasinteractive:
             matplotlib.interactive(True)
             install_repl_displayhook()
+        else:
+            matplotlib.interactive(False)
+            uninstall_repl_displayhook()
 
 
 class _IonContext:
@@ -408,6 +411,9 @@ def __exit__(self, exc_type, exc_value, traceback):
         if not self.wasinteractive:
             matplotlib.interactive(False)
             uninstall_repl_displayhook()
+        else:
+            matplotlib.interactive(True)
+            install_repl_displayhook()
 
 
 def ioff():
diff --git a/lib/matplotlib/tests/test_pyplot.py b/lib/matplotlib/tests/test_pyplot.py
@@ -110,22 +110,46 @@ def test_ion():
         assert mpl.is_interactive()
     assert mpl.is_interactive()
 
+
 def test_nested_ion_ioff():
+    # initial state is interactive
     plt.ion()
+
+    # mixed ioff/ion
     with plt.ioff():
         assert not mpl.is_interactive()
         with plt.ion():
             assert mpl.is_interactive()
         assert not mpl.is_interactive()
     assert mpl.is_interactive()
 
+    # redundant contexts
     with plt.ioff():
         with plt.ioff():
             assert not mpl.is_interactive()
     assert mpl.is_interactive()
 
+    with plt.ion():
+        plt.ioff()
+    assert mpl.is_interactive()
+
+    # initial state is not interactive
     plt.ioff()
+
+    # mixed ioff/ion
+    with plt.ion():
+        assert mpl.is_interactive()
+        with plt.ioff():
+            assert not mpl.is_interactive()
+        assert mpl.is_interactive()
+    assert not mpl.is_interactive()
+
+    # redunant contexts
     with plt.ion():
         with plt.ion():
             assert mpl.is_interactive()
     assert not mpl.is_interactive()
+
+    with plt.ioff():
+        plt.ion()
+    assert not mpl.is_interactive()
