add context manager functionality to ion and ioff

ianhi · ianhi · commit e1aa03e41462 · 2020-05-08T22:23:17.000-04:00
this enables with plt.ioff:
or with plt.ion:
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -333,16 +333,41 @@ def isinteractive():
     return matplotlib.is_interactive()
 
 
-def ioff():
-    """Turn the interactive mode off."""
-    matplotlib.interactive(False)
-    uninstall_repl_displayhook()
+class _ion_class():
+    def __call__(self):
+        """Turn the interactive mode on."""
+        matplotlib.interactive(True)
+        install_repl_displayhook()
 
+    def __enter__(self):
+        self.wasinteractive = isinteractive()
+        self.__call__()
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if not self.wasinteractive:
+            matplotlib.interactive(False)
+            uninstall_repl_displayhook()
+        del self.wasinteractive
+
+
+class _ioff_class():
+    def __call__(self):
+        """Turn the interactive mode off."""
+        matplotlib.interactive(False)
+        uninstall_repl_displayhook()
+
+    def __enter__(self):
+        self.wasinteractive = isinteractive()
+        self.__call__()
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        if self.wasinteractive:
+            matplotlib.interactive(True)
+            install_repl_displayhook()
+        del self.wasinteractive
 
-def ion():
-    """Turn the interactive mode on."""
-    matplotlib.interactive(True)
-    install_repl_displayhook()
+ioff = _ioff_class()
+ion = _ion_class()
 
 
 def pause(interval):
