ENH : add function to add displayhook for IPython

tacaswell · tacaswell · commit 96d3605b93c0 · 2015-05-25T00:49:40.000-04:00
- add `draw_all` function to pyplot
 - add `install_ipython_repl_hook` to pyplot which when
 run monkey-patches the ipython DisplayHook to call draw_all
 when ever the repl comes back.
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -22,6 +22,7 @@
 
 import sys
 import warnings
+import types
 
 import matplotlib
 import matplotlib.colorbar
@@ -108,6 +109,35 @@ def _backend_selection():
 _backend_mod, new_figure_manager, draw_if_interactive, _show = pylab_setup()
 
 
+def install_ipython_repl_hook():
+    try:
+        from IPython.core.displayhook import DisplayHook
+    except ImportError:
+        return
+    dh = sys.displayhook
+    # make sure we really have an IPython thing
+    if not isinstance(dh, DisplayHook):
+        return
+
+    orig_func = type(dh).finish_displayhook
+
+    def finish_displayhook(self):
+        draw_all()
+        return orig_func(self)
+
+    dh.finish_displayhook = types.MethodType(finish_displayhook, dh)
+
+
+def draw_all():
+    """
+    Redraw all figures registered with the pyplot
+    state machine.
+    """
+    for f_mgr in _pylab_helpers.Gcf.get_all_fig_managers():
+        # TODO add logic to check if figure is dirty
+        f_mgr.canvas.draw()
+
+
 @docstring.copy_dedent(Artist.findobj)
 def findobj(o=None, match=None, include_self=True):
     if o is None:
