Updated ElementBase.match and added unit tests.

2 files changed, 96 insertions, 24 deletions

diff --git a/sleekxmpp/xmlstream/stanzabase.py b/sleekxmpp/xmlstream/stanzabase.py
index 22853a9d..af74919e 100644
--- a/sleekxmpp/xmlstream/stanzabase.py
+++ b/sleekxmpp/xmlstream/stanzabase.py
@@ -454,7 +454,60 @@ class ElementBase(object):
                 # If we don't want to delete elements up the tree, stop
                 # after deleting the first level of elements.
                 return
+     
+    def match(self, xpath):
+        """
+        Compare a stanza object with an XPath expression. If the XPath matches
+        the contents of the stanza object, the match is successful.
+
+        The XPath expression may include checks for stanza attributes.
+        For example:
+            presence@show=xa@priority=2/status
+        Would match a presence stanza whose show value is set to 'xa', has a 
+        priority value of '2', and has a status element.
+
+        Arguments:
+            xpath -- The XPath expression to check against. It may be either a
+                     string or a list of element names with attribute checks.
+        """
+        if isinstance(xpath, str):
+            xpath = xpath.split('/')
+
+        # Extract the tag name and attribute checks for the first XPath node.
+        components = xpath[0].split('@')
+        tag = components[0]
+        attributes = components[1:]
         
+        if tag not in (self.name, self.plugins, self.plugin_attrib):
+            # The requested tag is not in this stanza, so no match.
+            return False
+
+        # Check the rest of the XPath against any substanzas.
+        matched_substanzas = False
+        for substanza in self.iterables:
+            if xpath[1:] == []:
+                break
+            matched_substanzas = substanza.match(xpath[1:])
+            if matched_substanzas: 
+                break
+
+        # Check attribute values.
+        for attribute in attributes:
+            name, value = attribute.split('=')
+            if self[name] != value:
+                return False
+
+        # Attempt to continue matching the XPath using the stanza's plugins.
+        if not matched_substanzas and len(xpath) > 1:
+            next_tag = xpath[1].split('@')[0]
+            if next_tag in self.plugins:
+                return self.plugins[next_tag].match(xpath[1:])
+            else:
+                return False
+
+        # Everything matched.
+        return True
+
     @property
     def attrib(self): #backwards compatibility
             return self
@@ -511,30 +564,6 @@ class ElementBase(object):
                     out.append('substanzas')
             return tuple(out)
 
-    def match(self, matchstring):
-            if isinstance(matchstring, str):
-                    nodes = matchstring.split('/')
-            else:
-                    nodes = matchstring
-            tagargs = nodes[0].split('@')
-            if tagargs[0] not in (self.plugins, self.plugin_attrib): return False
-            founditerable = False
-            for iterable in self.iterables:
-                    if nodes[1:] == []:
-                            break
-                    founditerable = iterable.match(nodes[1:])
-                    if founditerable: break;
-            for evals in tagargs[1:]:
-                    x,y = evals.split('=')
-                    if self[x] != y: return False
-            if not founditerable and len(nodes) > 1:
-                    next = nodes[1].split('@')[0]
-                    if next in self.plugins:
-                            return self.plugins[next].match(nodes[1:])
-                    else:
-                            return False
-            return True
-
     def find(self, xpath): # for backwards compatiblity, expose elementtree interface
             return self.xml.find(xpath)
 
diff --git a/tests/test_elementbase.py b/tests/test_elementbase.py
index 9628ea6c..d749f08e 100644
--- a/tests/test_elementbase.py
+++ b/tests/test_elementbase.py
@@ -428,5 +428,48 @@ class TestElementBase(SleekTest):
           </foo>
         """)
 
+    def testMatch(self):
+        """Test matching a stanza against an XPath expression."""
+
+        class TestSubStanza(ElementBase):
+            name = "sub"
+            namespace = "foo"
+            interfaces = set(('attrib',))
+
+        class TestStanza(ElementBase):
+            name = "foo"
+            namespace = "foo"
+            interfaces = set(('bar','baz'))
+            subitem = (TestSubStanza,)
+
+        class TestStanzaPlugin(ElementBase):
+            name = "plugin"
+            namespace = "foo"
+            interfaces = set(('attrib',))
+
+        registerStanzaPlugin(TestStanza, TestStanzaPlugin)
+
+        stanza = TestStanza()
+        self.failUnless(stanza.match("foo"), 
+            "Stanza did not match its own tag name.")
+
+        stanza['bar'] = 'a'
+        self.failUnless(stanza.match("foo@bar=a"),
+            "Stanza did not match its own name with attribute value check.")
+
+        stanza['baz'] = 'b'
+        self.failUnless(stanza.match("foo@bar=a@baz=b"),
+            "Stanza did not match its own name with multiple attributes.")
+
+        stanza['plugin']['attrib'] = 'c'
+        self.failUnless(stanza.match("foo/plugin@attrib=c"),
+            "Stanza did not match with plugin and attribute.")
+
+        substanza = TestSubStanza()
+        substanza['attrib'] = 'd'
+        stanza.append(substanza)
+        self.failUnless(stanza.match("foo/sub@attrib=d"),
+            "Stanza did not match with substanzas and attribute.")
+            
 
 suite = unittest.TestLoader().loadTestsFromTestCase(TestElementBase)