added findobj

jdh2358 · jdh2358 · commit 46ef01d32297 · 2008-07-03T14:30:13.000Z
svn path=/trunk/matplotlib/; revision=5709
diff --git a/CHANGELOG b/CHANGELOG
@@ -1,3 +1,6 @@
+2007-07-03 Implemented findobj method for artist and pyplot - see
+           examples/pylab_examples/findobj_demo.py - JDH
+
 2008-06-30 Another attempt to fix TextWithDash - DSD
 
 2008-06-30 Removed Qt4 NavigationToolbar2.destroy -- it appears to 
diff --git a/examples/pylab_examples/findobj_demo.py b/examples/pylab_examples/findobj_demo.py
@@ -0,0 +1,34 @@
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib.text as text
+
+a = np.arange(0,3,.02)
+b = np.arange(0,3,.02)
+c = np.exp(a)
+d = c[::-1]
+
+fig = plt.figure()
+ax = fig.add_subplot(111)
+plt.plot(a,c,'k--',a,d,'k:',a,c+d,'k')
+plt.legend(('Model length', 'Data length', 'Total message length'),
+           'upper center', shadow=True)
+plt.ylim([-1,20])
+plt.grid(False)
+plt.xlabel('Model complexity --->')
+plt.ylabel('Message length --->')
+plt.title('Minimum Message Length')
+
+# match on arbitrary function
+def myfunc(x):
+    return hasattr(x, 'set_color')
+
+for o in fig.findobj(myfunc):
+    o.set_color('blue')
+
+# match on class instances
+for o in fig.findobj(text.Text):
+    o.set_fontstyle('italic')
+
+
+
+plt.show()
diff --git a/lib/matplotlib/artist.py b/lib/matplotlib/artist.py
@@ -508,6 +508,46 @@ def set(self, **kwargs):
             ret.extend( [func(v)] )
         return ret
 
+    def findobj(self, match=None):
+        """
+        recursively find all :class:matplotlib.artist.Artist instances
+        contained in self
+
+        *match* can be
+
+          - None: return all objects contained in artist (including artist)
+
+          - function with signature ``boolean = match(artist)`` used to filter matches
+
+          - class instance: eg Line2D.  Only return artists of class type
+        """
+
+        if match is None: # always return True
+            def matchfunc(x): return True
+        elif cbook.issubclass_safe(match, Artist):
+            def matchfunc(x):
+                return isinstance(x, match)
+        elif callable(match):
+            matchfunc = match
+        else:
+            raise ValueError('match must be None, an matplotlib.artist.Artist subclass, or a callable')
+
+
+        artists = []
+        if hasattr(self, 'get_children'):
+            for c in self.get_children():
+                if matchfunc(c):
+                    artists.append(c)
+                artists.extend([thisc for thisc in c.findobj(matchfunc) if matchfunc(thisc)])
+        else:
+            if matchfunc(self):
+                artists.append(self)
+        return artists
+
+
+
+
+
 
 class ArtistInspector:
     """
@@ -690,6 +730,48 @@ def pprint_getters(self):
         return lines
 
 
+
+    def findobj(self, match=None):
+        """
+        recursively find all :class:matplotlib.artist.Artist instances
+        contained in self
+
+        if *match* is not None, it can be
+
+          - function with signature ``boolean = match(artist)``
+
+          - class instance: eg Line2D
+
+        used to filter matches
+        """
+
+        if match is None: # always return True
+            def matchfunc(x): return True
+        elif issubclass(match, Artist):
+            def matchfunc(x):
+                return isinstance(x, match)
+        elif callable(match):
+            matchfunc = func
+        else:
+            raise ValueError('match must be None, an matplotlib.artist.Artist subclass, or a callable')
+
+
+        artists = []
+        if hasattr(self, 'get_children'):
+            for c in self.get_children():
+                if matchfunc(c):
+                    artists.append(c)
+                artists.extend([thisc for thisc in c.findobj(matchfunc) if matchfunc(thisc)])
+        else:
+            if matchfunc(self):
+                artists.append(self)
+        return artists
+
+
+
+
+
+
 def getp(o, property=None):
     """
     Return the value of handle property.  property is an optional string
diff --git a/lib/matplotlib/cbook.py b/lib/matplotlib/cbook.py
@@ -886,6 +886,13 @@ def safezip(*args):
             raise ValueError(_safezip_msg % (Nx, i+1, len(arg)))
     return zip(*args)
 
+def issubclass_safe(x, klass):
+    'return issubclass(x, klass) and return False on a TypeError'
+
+    try:
+        return issubclass(x, klass)
+    except TypeError:
+        return False
 
 class MemoryMonitor:
     def __init__(self, nmax=20000):
diff --git a/lib/matplotlib/legend.py b/lib/matplotlib/legend.py
@@ -363,6 +363,12 @@ def draw_frame(self, b):
         'b is a boolean.  Set draw frame to b'
         self._drawFrame = b
 
+    def get_children(self):
+        children = []
+        children.extend(self.legendHandles)
+        children.extend(self.texts)
+        return children
+
     def get_frame(self):
         'return the Rectangle instance used to frame the legend'
         return self.legendPatch
diff --git a/lib/matplotlib/pylab.py b/lib/matplotlib/pylab.py
@@ -38,6 +38,7 @@
   figtext   - add text in figure coords
   figure   - create or change active figure
   fill     - make filled polygons
+  findobj  - recursively find all objects matching some criteria
   gca      - return the current axes
   gcf      - return the current figure
   gci      - get the current image, or None
diff --git a/lib/matplotlib/pyplot.py b/lib/matplotlib/pyplot.py
@@ -38,6 +38,29 @@
 from matplotlib.backends import pylab_setup
 new_figure_manager, draw_if_interactive, show = pylab_setup()
 
+
+
+def findobj(o=None, match=None):
+    """
+    recursively find all :class:matplotlib.artist.Artist instances
+    contained in artist instance *p*.  if *o* is None, use
+    current figure
+
+    *match* can be
+
+      - None: return all objects contained in artist (including artist)
+
+      - function with signature ``boolean = match(artist)`` used to filter matches
+
+      - class instance: eg Line2D.  Only return artists of class type
+
+    """
+
+    if o is None:
+        o = gcf()
+    return o.findobj(match)
+
+
 def switch_backend(newbackend):
     """
     Switch the default backend to newbackend.  This feature is
