Updated getStanzaValues and setStanzaValues with docs and unit tests.

2 files changed, 142 insertions, 33 deletions

diff --git a/sleekxmpp/xmlstream/stanzabase.py b/sleekxmpp/xmlstream/stanzabase.py
index aba63331..3e280d61 100644
--- a/sleekxmpp/xmlstream/stanzabase.py
+++ b/sleekxmpp/xmlstream/stanzabase.py
@@ -131,6 +131,58 @@ class ElementBase(object):
             self.plugins[attrib] = plugin_class(parent=self)
         return self
 
+    def getStanzaValues(self):
+        """
+        Return a dictionary of the stanza's interface values.
+
+        Stanza plugin values are included as nested dictionaries.
+        """
+        values = {}
+        for interface in self.interfaces:
+            values[interface] = self[interface]
+        for plugin, stanza in self.plugins.items():
+            values[plugin] = stanza.getStanzaValues()
+        if self.iterables:
+            iterables = []
+            for stanza in self.iterables:
+                iterables.append(stanza.getStanzaValues())
+                iterables[-1].update({
+                    '__childtag__': "{%s}%s" % (stanza.namespace, stanza.name)
+                    })
+            values['substanzas'] = iterables
+        return values
+
+    def setStanzaValues(self, values):
+        """
+        Set multiple stanza interface values using a dictionary.
+
+        Stanza plugin values may be set using nested dictionaries.
+
+        Arguments:
+            values -- A dictionary mapping stanza interface with values.
+                      Plugin interfaces may accept a nested dictionary that
+                      will be used recursively.
+        """
+        for interface, value in values.items():
+            if interface == 'substanzas':
+                for subdict in value:
+                    if '__childtag__' in subdict:
+                        for subclass in self.subitem:
+                            child_tag = "{%s}%s" % (subclass.namespace,
+                                                    subclass.name)
+                            if subdict['__childtag__'] == child_tag:
+                                sub = subclass(parent=self)
+                                sub.setStanzaValues(subdict)
+                                self.iterables.append(sub)
+                                break
+            elif interface in self.interfaces:
+                self[interface] = value
+            elif interface in self.plugin_attrib_map:
+                if interface not in self.plugins:
+                    self.initPlugin(interface)
+                self.plugins[interface].setStanzaValues(value)
+        return self
+
     @property
     def attrib(self): #backwards compatibility
             return self
@@ -313,39 +365,6 @@ class ElementBase(object):
                     if child.tag == name:
                             self.xml.remove(child)
 
-    def getStanzaValues(self):
-            out = {}
-            for interface in self.interfaces:
-                    out[interface] = self[interface]
-            for pluginkey in self.plugins:
-                    out[pluginkey] = self.plugins[pluginkey].getStanzaValues()
-            if self.iterables:
-                    iterables = []
-                    for stanza in self.iterables:
-                            iterables.append(stanza.getStanzaValues())
-                            iterables[-1].update({'__childtag__': "{%s}%s" % (stanza.namespace, stanza.name)})
-                    out['substanzas'] = iterables
-            return out
-
-    def setStanzaValues(self, attrib):
-            for interface in attrib:
-                    if interface == 'substanzas':
-                            for subdict in attrib['substanzas']:
-                                    if '__childtag__' in subdict:
-                                            for subclass in self.subitem:
-                                                    if subdict['__childtag__'] == "{%s}%s" % (subclass.namespace, subclass.name):
-                                                            sub = subclass(parent=self)
-                                                            sub.setStanzaValues(subdict)
-                                                            self.iterables.append(sub)
-                                                            break
-                    elif interface in self.interfaces:
-                            self[interface] = attrib[interface]
-                    elif interface in self.plugin_attrib_map and interface not in self.plugins:
-                            self.initPlugin(interface)
-                    if interface in self.plugins:
-                            self.plugins[interface].setStanzaValues(attrib[interface])
-            return self
-
     def appendxml(self, xml):
             self.xml.append(xml)
             return self
diff --git a/tests/test_elementbase.py b/tests/test_elementbase.py
index 1b018b44..99c397a5 100644
--- a/tests/test_elementbase.py
+++ b/tests/test_elementbase.py
@@ -19,5 +19,95 @@ class TestElementBase(SleekTest):
           </foo>
         """)
 
+    def testGetStanzaValues(self):
+        """Test getStanzaValues using plugins and substanzas."""
+
+        class TestStanzaPlugin(ElementBase):
+            name = "foo2"
+            namespace = "foo"
+            interfaces = set(('bar', 'baz'))
+            plugin_attrib = "foo2"
+
+        class TestSubStanza(ElementBase):
+            name = "subfoo"
+            namespace = "foo"
+            interfaces = set(('bar', 'baz'))
+
+        class TestStanza(ElementBase):
+            name = "foo"
+            namespace = "foo"
+            interfaces = set(('bar', 'baz'))
+            subitem = set((TestSubStanza,))
+
+        registerStanzaPlugin(TestStanza, TestStanzaPlugin)
+
+        stanza = TestStanza()
+        stanza['bar'] = 'a'
+        stanza['foo2']['baz'] = 'b'
+        substanza = TestSubStanza()
+        substanza['bar'] = 'c'
+        stanza.append(substanza)
+
+        values = stanza.getStanzaValues()
+        expected = {'bar': 'a',
+                    'baz': '',
+                    'foo2': {'bar': '',
+                             'baz': 'b'},
+                    'substanzas': [{'__childtag__': '{foo}subfoo',
+                                    'bar': 'c',
+                                    'baz': ''}]}
+        self.failUnless(values == expected,
+            "Unexpected stanza values:\n%s\n%s" % (str(expected), str(values)))
+
+
+    def testSetStanzaValues(self):
+        """Test using setStanzaValues with substanzas and plugins."""
+
+        class TestStanzaPlugin(ElementBase):
+            name = "pluginfoo"
+            namespace = "foo"
+            interfaces = set(('bar', 'baz'))
+            plugin_attrib = "plugin_foo"
+
+        class TestStanzaPlugin2(ElementBase):
+            name = "pluginfoo2"
+            namespace = "foo"
+            interfaces = set(('bar', 'baz'))
+            plugin_attrib = "plugin_foo2"
+
+        class TestSubStanza(ElementBase):
+            name = "subfoo"
+            namespace = "foo"
+            interfaces = set(('bar', 'baz'))
+
+        class TestStanza(ElementBase):
+            name = "foo"
+            namespace = "foo"
+            interfaces = set(('bar', 'baz'))
+            subitem = set((TestSubStanza,))
+
+        registerStanzaPlugin(TestStanza, TestStanzaPlugin)
+        registerStanzaPlugin(TestStanza, TestStanzaPlugin2)
+
+        stanza = TestStanza()
+        values = {'bar': 'a',
+                  'baz': '',
+                  'plugin_foo': {'bar': '',
+                                 'baz': 'b'},
+                  'plugin_foo2': {'bar': 'd',
+                                  'baz': 'e'},
+                  'substanzas': [{'__childtag__': '{foo}subfoo',
+                                  'bar': 'c',
+                                  'baz': ''}]}
+        stanza.setStanzaValues(values)
+
+        self.checkStanza(TestStanza, stanza, """
+          <foo xmlns="foo" bar="a">
+            <pluginfoo baz="b" />
+            <pluginfoo2 bar="d" baz="e" />
+            <subfoo bar="c" />
+          </foo>
+        """)
+
 
 suite = unittest.TestLoader().loadTestsFromTestCase(TestElementBase)