Also dynamically generate scalarmappable functions.

anntzer · anntzer · commit 2a1b4ee9c6e8 · 2018-01-09T17:46:43.000-08:00
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -316,26 +316,8 @@ def rcdefaults():
         draw_all()
 
 
-# The current "image" (ScalarMappable) is retrieved or set
-# only via the pyplot interface using the following two
-# functions:
-def gci():
-    """
-    Get the current colorable artist.  Specifically, returns the
-    current :class:`~matplotlib.cm.ScalarMappable` instance (image or
-    patch collection), or *None* if no images or patch collections
-    have been defined.  The commands :func:`~matplotlib.pyplot.imshow`
-    and :func:`~matplotlib.pyplot.figimage` create
-    :class:`~matplotlib.image.Image` instances, and the commands
-    :func:`~matplotlib.pyplot.pcolor` and
-    :func:`~matplotlib.pyplot.scatter` create
-    :class:`~matplotlib.collections.Collection` instances.  The
-    current image is an attribute of the current axes, or the nearest
-    earlier axes in the current figure that contains an image.
-    """
-    return gcf()._gci()
-
-
+# This function is not autogenerated because it is needed in when
+# autogenerating other functions, see `_post_processors`.
 def sci(im):
     """
     Set the current image.  This image will be the target of colormap
@@ -1991,23 +1973,6 @@ def getname_val(identifier):
         fig.autofmt_xdate()
 
 
-def _autogen_docstring(base):
-    """Autogenerated wrappers will get their docstring from a base function
-    with an addendum."""
-    msg = ''
-    addendum = docstring.Appender(msg, '\n\n')
-    return lambda func: addendum(docstring.copy_dedent(base)(func))
-
-
-# This function cannot be generated by boilerplate.py because it may
-# return an image or a line.
-@_autogen_docstring(Axes.spy)
-def spy(Z, precision=0, marker=None, markersize=None, aspect='equal', **kwargs):
-    ret = ax.spy(Z, precision, marker, markersize, aspect, **kwargs)
-    if isinstance(ret, cm.ScalarMappable):
-        sci(ret)
-    return ret
-
 # just to be safe.  Interactive mode can be turned on without
 # calling `plt.ion()` so register it again here.
 # This is safe because multiple calls to `install_repl_displayhook`
@@ -2022,11 +1987,14 @@ def _make_wrapper(obj_getter, cls, command):
     else:
         method_name = func_name = command
     method = getattr(cls, method_name)
+    post_processor = _post_processors.get(func_name, lambda obj: None)
 
     def func(*args, **kwargs):
         # We actually need to refetch the method here in case it has been
         # monkey-patched.
-        return getattr(obj_getter(), method_name)(*args, **kwargs)
+        ret = getattr(obj_getter(), method_name)(*args, **kwargs)
+        post_processor(ret)
+        return ret
 
     if six.PY2:
         func.__name__ = func_name.encode("ascii")
@@ -2047,7 +2015,7 @@ def func(*args, **kwargs):
     "clf", "figimage", "gca", "ginput", "savefig", "subplots_adjust",
     "suptitle", "tight_layout", "waitforbuttonpress",
     # Renamed commands.
-    "legend:figlegend", "text:figtext"]
+    "_gci:gci", "legend:figlegend", "text:figtext"]
 _axes_commands = [
     "acorr", "angle_spectrum", "annotate", "arrow", "autoscale", "axhline",
     "axhspan", "axis", "axvline", "axvspan", "bar", "barbs", "barh", "boxplot",
@@ -2063,7 +2031,26 @@ def func(*args, **kwargs):
     # Renamed commands.
     "set_title:title", "set_xlabel:xlabel", "set_xscale:xscale",
     "set_ylabel:ylabel", "set_yscale:yscale"]
-
+_post_processors = {
+    # For the following functions, an additional callable will be called with
+    # the return value as argument.
+    "hexbin": sci,
+    "imshow": sci,
+    "pcolor": sci,
+    "pcolormesh": sci,
+    "quiver": sci,
+    "scatter": sci,
+    "tripcolor": sci,
+    "contour": lambda ret: sci(ret) if ret._A is not None else None,
+    "contourf": lambda ret: sci(ret) if ret._A is not None else None,
+    "tricontour": lambda ret: sci(ret) if ret._A is not None else None,
+    "tricontourf": lambda ret: sci(ret) if ret._A is not None else None,
+    "hist2d": lambda ret: sci(ret[-1]),
+    "specgram": lambda ret: sci(ret[-1]),
+    "streamplot": lambda ret: sci(ret.lines),
+    "spy": (
+        lambda ret: sci(ret) if isinstance(ret, cm.ScalarMappable) else None)
+}
 
 for _command in _canvas_commands:
     _make_wrapper(
