Updated ElementBase.__setitem__ and added unit tests.

2 files changed, 84 insertions, 19 deletions

diff --git a/sleekxmpp/xmlstream/stanzabase.py b/sleekxmpp/xmlstream/stanzabase.py
index 962cf8ed..83a8ddfc 100644
--- a/sleekxmpp/xmlstream/stanzabase.py
+++ b/sleekxmpp/xmlstream/stanzabase.py
@@ -225,6 +225,50 @@ class ElementBase(object):
         else:
             return ''
 
+    def __setitem__(self, attrib, value):
+        """
+        Set the value of a stanza interface using dictionary-like syntax.
+
+        Example:
+            >>> msg['body'] = "Hi!"
+            >>> msg['body']
+            'Hi!'
+
+        Stanza interfaces are typically mapped directly to the underlying XML
+        object, but can be overridden by the presence of a setAttrib method
+        (or setFoo where the interface is named foo, etc).
+
+        The effect of interface value assignment for an interface
+        named 'foo' will be one of:
+            1. Delete the interface's contents if the value is None.
+            2. Call setFoo, if it exists.
+            3. Set the text of a foo element, if foo is in sub_interfaces.
+            4. Set the value of a top level XML attribute name foo.
+            5. Attempt to pass value to a plugin named foo using the plugin's
+               foo interface.
+            6. Do nothing.
+
+        Arguments:
+            attrib -- The name of the stanza interface to modify.
+            value  -- The new value of the stanza interface.
+        """
+        if attrib in self.interfaces:
+            if value is not None:
+                if hasattr(self, "set%s" % attrib.title()):
+                    getattr(self, "set%s" % attrib.title())(value,)
+                else:
+                    if attrib in self.sub_interfaces:
+                        return self._setSubText(attrib, text=value)
+                    else:
+                        self._setAttr(attrib, value)
+            else:
+                self.__delitem__(attrib)
+        elif attrib in self.plugin_attrib_map:
+            if attrib not in self.plugins:
+                self.initPlugin(attrib)
+            self.plugins[attrib][attrib] = value
+        return self
+
     @property
     def attrib(self): #backwards compatibility
             return self
@@ -308,24 +352,6 @@ class ElementBase(object):
     def findall(self, xpath):
             return self.xml.findall(xpath)
 
-    def __setitem__(self, attrib, value):
-            if attrib in self.interfaces:
-                    if value is not None:
-                            if hasattr(self, "set%s" % attrib.title()):
-                                    getattr(self, "set%s" % attrib.title())(value,)
-                            else:
-                                    if attrib in self.sub_interfaces:
-                                            return self._setSubText(attrib, text=value)
-                                    else:
-                                            self._setAttr(attrib, value)
-                    else:
-                            self.__delitem__(attrib)
-            elif attrib in self.plugin_attrib_map:
-                    if attrib not in self.plugins: self.initPlugin(attrib)
-                    self.initPlugin(attrib)
-                    self.plugins[attrib][attrib] = value
-            return self
-
     def __delitem__(self, attrib):
             if attrib.lower() in self.interfaces:
                     if hasattr(self, "del%s" % attrib.title()):
diff --git a/tests/test_elementbase.py b/tests/test_elementbase.py
index d6fd457c..95502f54 100644
--- a/tests/test_elementbase.py
+++ b/tests/test_elementbase.py
@@ -115,9 +115,12 @@ class TestElementBase(SleekTest):
         class TestStanza(ElementBase):
             name = "foo"
             namespace = "foo"
-            interfaces = set(('bar', 'baz'))
+            interfaces = set(('bar', 'baz', 'qux'))
             sub_interfaces = set(('baz',))
 
+            def getQux(self):
+              return 'qux'
+
         class TestStanzaPlugin(ElementBase):
             name = "foobar"
             namespace = "foo"
@@ -132,12 +135,14 @@ class TestElementBase(SleekTest):
         stanza.append(substanza)
         stanza.setStanzaValues({'bar': 'a',
                                 'baz': 'b',
+                                'qux': 42,
                                 'foobar': {'fizz': 'c'}})
 
         # Test non-plugin interfaces
         expected = {'substanzas': [substanza],
                     'bar': 'a',
                     'baz': 'b',
+                    'qux': 'qux',
                     'meh': ''}
         for interface, value in expected.items():
             result = stanza[interface]
@@ -150,5 +155,39 @@ class TestElementBase(SleekTest):
         self.failUnless(stanza['foobar']['fizz'] == 'c',
                         "Incorrect plugin subvalue result.")
 
+    def testSetItem(self):
+        """Test assigning to stanza interfaces."""
+
+        class TestStanza(ElementBase):
+            name = "foo"
+            namespace = "foo"
+            interfaces = set(('bar', 'baz', 'qux'))
+            sub_interfaces = set(('baz',))
+
+            def setQux(self, value):
+                pass
+
+        class TestStanzaPlugin(ElementBase):
+            name = "foobar"
+            namespace = "foo"
+            plugin_attrib = "foobar"
+            interfaces = set(('foobar',))
+
+        registerStanzaPlugin(TestStanza, TestStanzaPlugin)
+
+        stanza = TestStanza()
+
+        stanza['bar'] = 'attribute!'
+        stanza['baz'] = 'element!'
+        stanza['qux'] = 'overridden'
+        stanza['foobar'] = 'plugin'
+
+        self.checkStanza(TestStanza, stanza, """
+          <foo xmlns="foo" bar="attribute!">
+            <baz>element!</baz>
+            <foobar foobar="plugin" />
+          </foo>
+        """)
+
 
 suite = unittest.TestLoader().loadTestsFromTestCase(TestElementBase)