Updated ElementBase._delSub and added unit tests.

_delSub can now accept a path and will optionally remove any empty parent elements after deleting the target elements.

2 files changed, 114 insertions, 7 deletions

diff --git a/sleekxmpp/xmlstream/stanzabase.py b/sleekxmpp/xmlstream/stanzabase.py
index bccb1fb0..22853a9d 100644
--- a/sleekxmpp/xmlstream/stanzabase.py
+++ b/sleekxmpp/xmlstream/stanzabase.py
@@ -417,6 +417,44 @@ class ElementBase(object):
         element.text = text
         return element 
 
+    def _delSub(self, name, all=False):
+        """
+        Remove sub elements that match the given name or XPath.
+
+        If the element is in a path, then any parent elements that become
+        empty after deleting the element may also be deleted if requested
+        by setting all=True.
+
+        Arguments:
+            name -- The name or XPath expression for the element(s) to remove.
+            all  -- If True, remove all empty elements in the path to the 
+                    deleted element. Defaults to False.
+        """
+        name = self._fix_ns(name)
+        path = name.split("/")
+        original_target = path[-1]
+
+        for level, _ in enumerate(path):
+            # Generate the paths to the target elements and their parent.
+            element_path = "/".join(path[:len(path) - level])
+            parent_path = "/".join(path[:len(path) - level - 1])
+
+            elements = self.xml.findall(element_path)
+            parent = self.xml.find(parent_path)
+
+            if elements:
+                if parent is None:
+                    parent = self.xml
+                for element in elements:
+                    if element.tag == original_target or not element.getchildren():
+                        # Only delete the originally requested elements, and any
+                        # parent elements that have become empty.
+                        parent.remove(element)
+            if not all:
+                # If we don't want to delete elements up the tree, stop
+                # after deleting the first level of elements.
+                return
+        
     @property
     def attrib(self): #backwards compatibility
             return self
@@ -512,13 +550,6 @@ class ElementBase(object):
                             return False
             return True
 
-    def _delSub(self, name):
-            if '}' not in name:
-                    name = "{%s}%s" % (self.namespace, name)
-            for child in self.xml.getchildren():
-                    if child.tag == name:
-                            self.xml.remove(child)
-
     def appendxml(self, xml):
             self.xml.append(xml)
             return self
diff --git a/tests/test_elementbase.py b/tests/test_elementbase.py
index 2b61489a..9628ea6c 100644
--- a/tests/test_elementbase.py
+++ b/tests/test_elementbase.py
@@ -352,5 +352,81 @@ class TestElementBase(SleekTest):
           </foo>
         """)
 
+    def testDelSub(self):
+        """Test removing sub elements."""
+
+        class TestStanza(ElementBase):
+            name = "foo"
+            namespace = "foo"
+            interfaces = set(('bar', 'baz'))
+
+            def setBar(self, value):
+                self._setSubText("path/to/only/bar", value);
+
+            def getBar(self):
+                return self._getSubText("path/to/only/bar")
+
+            def delBar(self):
+                self._delSub("path/to/only/bar")
+
+            def setBaz(self, value):
+                self._setSubText("path/to/just/baz", value);
+
+            def getBaz(self):
+                return self._getSubText("path/to/just/baz")
+
+            def delBaz(self):
+                self._delSub("path/to/just/baz")
+
+        stanza = TestStanza()
+        stanza['bar'] = 'a'
+        stanza['baz'] = 'b'
+
+        self.checkStanza(TestStanza, stanza, """
+          <foo xmlns="foo">
+            <path>
+              <to>
+                <only>
+                  <bar>a</bar>
+                </only>
+                <just>
+                  <baz>b</baz>
+                </just>
+              </to>
+            </path>
+          </foo>
+        """)
+
+        del stanza['bar']
+        del stanza['baz']
+
+        self.checkStanza(TestStanza, stanza, """
+          <foo xmlns="foo">
+            <path>
+              <to>
+                <only />
+                <just />
+              </to>
+            </path>
+          </foo>
+        """, use_values=False)
+
+        stanza['bar'] = 'a'
+        stanza['baz'] = 'b'
+
+        stanza._delSub('path/to/only/bar', all=True)
+
+        self.checkStanza(TestStanza, stanza, """
+          <foo xmlns="foo">
+            <path>
+              <to>
+                <just>
+                  <baz>b</baz>
+                </just>
+              </to>
+            </path>
+          </foo>
+        """)
+
 
 suite = unittest.TestLoader().loadTestsFromTestCase(TestElementBase)