Merge pull request #17371 from ianhi/context-manager

timhoffm · web-flow · commit 89d13c30fe2b · 2020-08-12T09:13:39.000+02:00
add context manager functionality to ion and ioff
diff --git a/doc/api/next_api_changes/behavior/17371-IHI.rst b/doc/api/next_api_changes/behavior/17371-IHI.rst
@@ -0,0 +1,7 @@
+ioff and ion can be used as context managers
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+`.pyplot.ion` and `.pyplot.ioff` may now be used as context managers to create
+a context with interactive mode on, or off respectively. The old behavior of
+calling these functions is maintained. To use the new functionality
+call as ``with plt.ioff():``
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -364,9 +364,61 @@ def isinteractive():
     return matplotlib.is_interactive()
 
 
+class _IoffContext:
+    """
+    Context manager for `.ioff`.
+
+    The state is changed in ``__init__()`` instead of ``__enter__()``. The
+    latter is a no-op. This allows using `.ioff` both as a function and
+    as a context.
+    """
+
+    def __init__(self):
+        self.wasinteractive = isinteractive()
+        matplotlib.interactive(False)
+        uninstall_repl_displayhook()
+
+    def __enter__(self):
+        pass
+
+    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:
+    """
+    Context manager for `.ion`.
+
+    The state is changed in ``__init__()`` instead of ``__enter__()``. The
+    latter is a no-op. This allows using `.ion` both as a function and
+    as a context.
+    """
+
+    def __init__(self):
+        self.wasinteractive = isinteractive()
+        matplotlib.interactive(True)
+        install_repl_displayhook()
+
+    def __enter__(self):
+        pass
+
+    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():
     """
-    Turn the interactive mode off.
+    Turn interactive mode off.
 
     See Also
     --------
@@ -375,14 +427,33 @@ def ioff():
 
     show : show windows (and maybe block)
     pause : show windows, run GUI event loop, and block for a time
+
+    Notes
+    -----
+    For a temporary change, this can be used as a context manager::
+
+        # if interactive mode is on
+        # then figures will be shown on creation
+        plt.ion()
+        # This figure will be shown immediately
+        fig = plt.figure()
+
+        with plt.ioff():
+            # interactive mode will be off
+            # figures will not automatically be shown
+            fig2 = plt.figure()
+            # ...
+
+    To enable usage as a context manager, this function returns an
+    ``_IoffContext`` object. The return value is not intended to be stored
+    or accessed by the user.
     """
-    matplotlib.interactive(False)
-    uninstall_repl_displayhook()
+    return _IoffContext()
 
 
 def ion():
     """
-    Turn the interactive mode on.
+    Turn interactive mode on.
 
     See Also
     --------
@@ -391,9 +462,28 @@ def ion():
 
     show : show windows (and maybe block)
     pause : show windows, run GUI event loop, and block for a time
+
+    Notes
+    -----
+    For a temporary change, this can be used as a context manager::
+
+        # if interactive mode is off
+        # then figures will not be shown on creation
+        plt.ioff()
+        # This figure will not be shown immediately
+        fig = plt.figure()
+
+        with plt.ion():
+            # interactive mode will be on
+            # figures will automatically be shown
+            fig2 = plt.figure()
+            # ...
+
+    To enable usage as a context manager, this function returns an
+    ``_IonContext`` object. The return value is not intended to be stored
+    or accessed by the user.
     """
-    matplotlib.interactive(True)
-    install_repl_displayhook()
+    return _IonContext()
 
 
 def pause(interval):
diff --git a/lib/matplotlib/tests/test_pyplot.py b/lib/matplotlib/tests/test_pyplot.py
@@ -81,3 +81,75 @@ def test_nrows_error():
         plt.subplot(nrows=1)
     with pytest.raises(TypeError):
         plt.subplot(ncols=1)
+
+
+def test_ioff():
+    plt.ion()
+    assert mpl.is_interactive()
+    with plt.ioff():
+        assert not mpl.is_interactive()
+    assert mpl.is_interactive()
+
+    plt.ioff()
+    assert not mpl.is_interactive()
+    with plt.ioff():
+        assert not mpl.is_interactive()
+    assert not mpl.is_interactive()
+
+
+def test_ion():
+    plt.ioff()
+    assert not mpl.is_interactive()
+    with plt.ion():
+        assert mpl.is_interactive()
+    assert not mpl.is_interactive()
+
+    plt.ion()
+    assert mpl.is_interactive()
+    with plt.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()
