add context manager functionality to ion and ioff

ianhi · ianhi · commit 8c3e60c76658 · 2020-07-29T00:20:18.000-04:00
this enables usage of plt.ioff as a context manager
to allow creating figures without the figure automatically displaying
in environments such as jupyter.

Usage:
with plt.ioff:
   plt.figure()
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -364,36 +364,62 @@ def isinteractive():
     return matplotlib.is_interactive()
 
 
-def ioff():
-    """
-    Turn the interactive mode off.
+class _ioff_class():
+    def __call__(self):
+        """
+        Turn the interactive mode off.
+
+        See Also
+        --------
+        ion : enable interactive mode
+        isinteractive : query current state
+
+        show : show windows (and maybe block)
+        pause : show windows, run GUI event loop, and block for a time
+        """
+        matplotlib.interactive(False)
+        uninstall_repl_displayhook()
 
-    See Also
-    --------
-    ion : enable interactive mode
-    isinteractive : query current state
+    def __enter__(self):
+        self.wasinteractive = isinteractive()
+        self.__call__()
 
-    show : show windows (and maybe block)
-    pause : show windows, run GUI event loop, and block for a time
-    """
-    matplotlib.interactive(False)
-    uninstall_repl_displayhook()
+    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.
+class _ion_class():
+    def __call__(self):
+        """
+        Turn the interactive mode on.
 
-    See Also
-    --------
-    ioff : disable interactive mode
-    isinteractive : query current state
+        See Also
+        --------
+        ioff : disable interactive mode
+        isinteractive : query current state
 
-    show : show windows (and maybe block)
-    pause : show windows, run GUI event loop, and block for a time
-    """
-    matplotlib.interactive(True)
-    install_repl_displayhook()
+        show : show windows (and maybe block)
+        pause : show windows, run GUI event loop, and block for a time
+        """
+        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
+
+
+ioff = _ioff_class()
+ion = _ion_class()
 
 
 def pause(interval):
