Tidy up and add tests for multi_attrib plugins.

2 files changed, 191 insertions, 5 deletions

diff --git a/sleekxmpp/xmlstream/stanzabase.py b/sleekxmpp/xmlstream/stanzabase.py
index 3b148446..5d572437 100644
--- a/sleekxmpp/xmlstream/stanzabase.py
+++ b/sleekxmpp/xmlstream/stanzabase.py
@@ -36,6 +36,17 @@ def register_stanza_plugin(stanza, plugin, iterable=False, overrides=False):
     >>> from sleekxmpp.xmlstream import register_stanza_plugin
     >>> register_stanza_plugin(Iq, CustomStanza)
 
+    Plugin stanzas marked as iterable will be included in the list of
+    substanzas for the parent, using ``parent['substanzas']``. If the
+    attribute ``plugin_multi_attrib`` was defined for the plugin, then
+    the substanza set can be filtered to only instances of the plugin
+    class. For example, given a plugin class ``Foo`` with 
+    ``plugin_multi_attrib = 'foos'`` then::
+
+        parent['foos']
+
+    would return a collection of all ``Foo`` substanzas.
+
     :param class stanza: The class of the parent stanza.
     :param class plugin: The class of the plugin stanza.
     :param bool iterable: Indicates if the plugin stanza should be
@@ -67,9 +78,9 @@ def register_stanza_plugin(stanza, plugin, iterable=False, overrides=False):
 
     if iterable:
         stanza.plugin_iterables.add(plugin)
-    if iterable and hasattr(plugin, 'multi_attrib'):
-        multiplugin = multifactory(plugin, plugin.multi_attrib)
-        register_stanza_plugin(stanza, multiplugin)
+        if plugin.plugin_multi_attrib:
+            multiplugin = multifactory(plugin, plugin.plugin_multi_attrib)
+            register_stanza_plugin(stanza, multiplugin)
     if overrides:
         for interface in plugin.overrides:
             stanza.plugin_overrides[interface] = plugin.plugin_attrib
@@ -93,7 +104,7 @@ def multifactory(stanza, plugin_attrib):
     def get_multi(self):
         parent = self.parent()
         res = filter(lambda sub: isinstance(sub, self._multistanza), parent)
-        return tuple(res)
+        return list(res)
 
     def set_multi(self, val):
         parent = self.parent()
@@ -104,7 +115,7 @@ def multifactory(stanza, plugin_attrib):
     def del_multi(self):
         parent = self.parent()
         res = filter(lambda sub: isinstance(sub, self._multistanza), parent)
-        for stanza in res:
+        for stanza in list(res):
             parent.iterables.remove(stanza)
             parent.xml.remove(stanza.xml)
 
@@ -258,6 +269,16 @@ class ElementBase(object):
     #:     msg['foo']['an_interface_from_the_foo_plugin']
     plugin_attrib = 'plugin'
 
+    #: For :class:`ElementBase` subclasses that are intended to be an
+    #: iterable group of items, the ``plugin_multi_attrib`` value defines
+    #: an interface for the parent stanza which returns the entire group
+    #: of matching substanzas. So the following are equivalent::
+    #:
+    #:     # Given stanza class Foo, with plugin_multi_attrib = 'foos'
+    #:     parent['foos']
+    #:     filter(isinstance(item, Foo), parent['substanzas'])
+    plugin_multi_attrib = ''
+
     #: The set of keys that the stanza provides for accessing and
     #: manipulating the underlying XML object. This set may be augmented
     #: with the :attr:`plugin_attrib` value of any registered
@@ -445,6 +466,8 @@ class ElementBase(object):
             self.plugins[attrib] = plugin
             if plugin_class in self.plugin_iterables:
                 self.iterables.append(plugin)
+                if plugin_class.plugin_multi_attrib:
+                    self.init_plugin(plugin_class.plugin_multi_attrib)
         return self
 
     def _get_stanza_values(self):
@@ -995,6 +1018,9 @@ class ElementBase(object):
                 raise TypeError
         self.xml.append(item.xml)
         self.iterables.append(item)
+        if item.__class__ in self.plugin_iterables:
+            if item.__class__.plugin_multi_attrib:
+                self.init_plugin(item.__class__.plugin_multi_attrib)
         return self
 
     def appendxml(self, xml):
diff --git a/tests/test_stanza_element.py b/tests/test_stanza_element.py
index 60ccb7c9..09093003 100644
--- a/tests/test_stanza_element.py
+++ b/tests/test_stanza_element.py
@@ -774,5 +774,165 @@ class TestElementBase(SleekTest):
           <foo xmlns="foo" />
         """)
 
+    def testGetMultiAttrib(self):
+        """Test retrieving multi_attrib substanzas."""
+
+        class TestStanza(ElementBase):
+            name = 'foo'
+            namespace = 'foo'
+            interfaces = set()
+
+        class TestMultiStanza1(ElementBase):
+            name = 'bar'
+            namespace = 'bar'
+            plugin_attrib = name
+            plugin_multi_attrib = 'bars'
+
+        class TestMultiStanza2(ElementBase):
+            name = 'baz'
+            namespace = 'baz'
+            plugin_attrib = name
+            plugin_multi_attrib = 'bazs'
+           
+        register_stanza_plugin(TestStanza, TestMultiStanza1, iterable=True)
+        register_stanza_plugin(TestStanza, TestMultiStanza2, iterable=True)
+
+        stanza = TestStanza()
+        stanza.append(TestMultiStanza1())
+        stanza.append(TestMultiStanza2())
+        stanza.append(TestMultiStanza1())
+        stanza.append(TestMultiStanza2())
+
+        self.check(stanza, """
+          <foo xmlns="foo">
+            <bar xmlns="bar" />
+            <baz xmlns="baz" />
+            <bar xmlns="bar" />
+            <baz xmlns="baz" />
+          </foo>
+        """, use_values=False)
+
+        bars = stanza['bars']
+        bazs = stanza['bazs']
+
+        for bar in bars:
+            self.check(bar, """
+              <bar xmlns="bar" />
+            """)
+
+        for baz in bazs:
+            self.check(baz, """
+              <baz xmlns="baz" />
+            """)
+
+        self.assertEqual(len(bars), 2, 
+                "Wrong number of <bar /> stanzas: %s" % len(bars))
+        self.assertEqual(len(bazs), 2, 
+                "Wrong number of <baz /> stanzas: %s" % len(bazs))
+
+    def testSetMultiAttrib(self):
+        """Test setting multi_attrib substanzas."""
+
+        class TestStanza(ElementBase):
+            name = 'foo'
+            namespace = 'foo'
+            interfaces = set()
+
+        class TestMultiStanza1(ElementBase):
+            name = 'bar'
+            namespace = 'bar'
+            plugin_attrib = name
+            plugin_multi_attrib = 'bars'
+
+        class TestMultiStanza2(ElementBase):
+            name = 'baz'
+            namespace = 'baz'
+            plugin_attrib = name
+            plugin_multi_attrib = 'bazs'
+           
+        register_stanza_plugin(TestStanza, TestMultiStanza1, iterable=True)
+        register_stanza_plugin(TestStanza, TestMultiStanza2, iterable=True)
+
+        stanza = TestStanza()
+        stanza['bars'] = [TestMultiStanza1(), TestMultiStanza1()]
+        stanza['bazs'] = [TestMultiStanza2(), TestMultiStanza2()]
+
+        self.check(stanza, """
+          <foo xmlns="foo">
+            <bar xmlns="bar" />
+            <bar xmlns="bar" />
+            <baz xmlns="baz" />
+            <baz xmlns="baz" />
+          </foo>
+        """, use_values=False)
+
+        self.assertEqual(len(stanza['substanzas']), 4,
+            "Wrong number of substanzas: %s" % len(stanza['substanzas']))
+
+        stanza['bars'] = [TestMultiStanza1()]
+
+        self.check(stanza, """
+          <foo xmlns="foo">
+            <baz xmlns="baz" />
+            <baz xmlns="baz" />
+            <bar xmlns="bar" />
+          </foo>
+        """, use_values=False)
+
+        self.assertEqual(len(stanza['substanzas']), 3,
+            "Wrong number of substanzas: %s" % len(stanza['substanzas']))
+
+
+    def testDelMultiAttrib(self):
+        """Test deleting multi_attrib substanzas."""
+
+        class TestStanza(ElementBase):
+            name = 'foo'
+            namespace = 'foo'
+            interfaces = set()
+
+        class TestMultiStanza1(ElementBase):
+            name = 'bar'
+            namespace = 'bar'
+            plugin_attrib = name
+            plugin_multi_attrib = 'bars'
+
+        class TestMultiStanza2(ElementBase):
+            name = 'baz'
+            namespace = 'baz'
+            plugin_attrib = name
+            plugin_multi_attrib = 'bazs'
+           
+        register_stanza_plugin(TestStanza, TestMultiStanza1, iterable=True)
+        register_stanza_plugin(TestStanza, TestMultiStanza2, iterable=True)
+
+        stanza = TestStanza()
+        bars = [TestMultiStanza1(), TestMultiStanza1()]
+        bazs = [TestMultiStanza2(), TestMultiStanza2()]
+
+        stanza['bars'] = bars
+        stanza['bazs'] = bazs
+
+        self.check(stanza, """
+          <foo xmlns="foo">
+            <bar xmlns="bar" />
+            <bar xmlns="bar" />
+            <baz xmlns="baz" />
+            <baz xmlns="baz" />
+          </foo>
+        """, use_values=False)
+
+        del stanza['bars']
+
+        self.check(stanza, """
+          <foo xmlns="foo">
+            <baz xmlns="baz" />
+            <baz xmlns="baz" />
+          </foo>
+        """, use_values=False)
+
+        self.assertEqual(len(stanza['substanzas']), 2,
+            "Wrong number of substanzas: %s" % len(stanza['substanzas']))
+
 
 suite = unittest.TestLoader().loadTestsFromTestCase(TestElementBase)