Return container object containing lines and arrows.

tonysyu · tonysyu · commit 57a7d3de8985 · 2012-09-02T09:16:37.000-04:00
diff --git a/lib/matplotlib/axes.py b/lib/matplotlib/axes.py
@@ -6601,17 +6601,17 @@ def streamplot(self, x, y, u, v, density=1, linewidth=None, color=None,
                    cmap=None, norm=None, arrowsize=1, arrowstyle='-|>',
                    minlength=0.1, transform=None):
         if not self._hold: self.cla()
-        lines = mstream.streamplot(self, x, y, u, v,
-                                   density=density,
-                                   linewidth=linewidth,
-                                   color=color,
-                                   cmap=cmap,
-                                   norm=norm,
-                                   arrowsize=arrowsize,
-                                   arrowstyle=arrowstyle,
-                                   minlength=minlength,
-                                   transform=transform)
-        return lines
+        stream_container = mstream.streamplot(self, x, y, u, v,
+                                              density=density,
+                                              linewidth=linewidth,
+                                              color=color,
+                                              cmap=cmap,
+                                              norm=norm,
+                                              arrowsize=arrowsize,
+                                              arrowstyle=arrowstyle,
+                                              minlength=minlength,
+                                              transform=transform)
+        return stream_container
     streamplot.__doc__ = mstream.streamplot.__doc__
 
     @docstring.dedent_interpd
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -3055,7 +3055,7 @@ def streamplot(x, y, u, v, density=1, linewidth=None, color=None, cmap=None,
         draw_if_interactive()
     finally:
         ax.hold(washold)
-    sci(ret[0])
+    sci(ret.lines)
     return ret
 
 # This function was autogenerated by boilerplate.py.  Do not edit as
diff --git a/lib/matplotlib/streamplot.py b/lib/matplotlib/streamplot.py
@@ -8,6 +8,8 @@
 import matplotlib.colors as mcolors
 import matplotlib.collections as mcollections
 import matplotlib.patches as patches
+import matplotlib.container as container
+
 
 __all__ = ['streamplot']
 
@@ -46,11 +48,13 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
     *minlength* : float
         Minimum length of streamline in axes coordinates.
 
-    Returns *streamlines* : :class:`~matplotlib.collections.LineCollection`
-        Line collection with all streamlines as a series of line segments.
-        Currently, there is no way to differentiate between line segments
-        on different streamlines (other than manually checking that segments
-        are connected).
+    Returns
+    -------
+    *stream_container* : StreamplotContainer
+        Container object with attributes
+            lines : `matplotlib.collections.LineCollection` of streamlines
+            arrows : collection of `matplotlib.patches.FancyArrowPatch` objects
+                repesenting arrows half-way along stream lines.
     """
     grid = Grid(x, y)
     mask = StreamMask(density)
@@ -154,8 +158,20 @@ def streamplot(axes, x, y, u, v, density=1, linewidth=None, color=None,
     axes.update_datalim(((x.min(), y.min()), (x.max(), y.max())))
     axes.autoscale_view(tight=True)
 
-    arrow_collection = matplotlib.collections.PatchCollection(arrows)
-    return lc, arrow_collection
+    ac = matplotlib.collections.PatchCollection(arrows)
+    stream_container = StreamplotContainer(lc, arrows=ac)
+    return stream_container
+
+
+class StreamplotContainer(container.Container):
+
+    def __new__(cls, *kl, **kwargs):
+        return tuple.__new__(cls)
+
+    def __init__(self, lines, arrows=None, **kwargs):
+        self.lines = lines
+        self.arrows = arrows
+        container.Container.__init__(self, lines, **kwargs)
 
 
 # Coordinate definitions
