XMPPError exceptions can keep a stanza's contents.

This allows exceptions to include the original
content of a stanza in the error response by including
the parameter clear=False when raising the exception.

7 files changed, 71 insertions, 12 deletions

diff --git a/sleekxmpp/exceptions.py b/sleekxmpp/exceptions.py
index d3988b4a..4727f0c6 100644
--- a/sleekxmpp/exceptions.py
+++ b/sleekxmpp/exceptions.py
@@ -21,7 +21,8 @@ class XMPPError(Exception):
     """
 
     def __init__(self, condition='undefined-condition', text=None, etype=None,
-                 extension=None, extension_ns=None, extension_args=None):
+                 extension=None, extension_ns=None, extension_args=None,
+                 clear=True):
         """
         Create a new XMPPError exception.
 
@@ -37,6 +38,9 @@ class XMPPError(Exception):
             extension_args -- Content and attributes for the extension
                               element. Same as the additional arguments to
                               the ET.Element constructor.
+            clear          -- Indicates if the stanza's contents should be
+                              removed before replying with an error.
+                              Defaults to True.
         """
         if extension_args is None:
             extension_args = {}
@@ -44,6 +48,7 @@ class XMPPError(Exception):
         self.condition = condition
         self.text = text
         self.etype = etype
+        self.clear = clear
         self.extension = extension
         self.extension_ns = extension_ns
         self.extension_args = extension_args
diff --git a/sleekxmpp/stanza/iq.py b/sleekxmpp/stanza/iq.py
index c6aa64d0..841d282c 100644
--- a/sleekxmpp/stanza/iq.py
+++ b/sleekxmpp/stanza/iq.py
@@ -144,7 +144,7 @@ class Iq(RootStanza):
                 self.xml.remove(child)
         return self
 
-    def reply(self):
+    def reply(self, clear=True):
         """
         Send a reply <iq> stanza.
 
@@ -152,9 +152,13 @@ class Iq(RootStanza):
 
         Sets the 'type' to 'result' in addition to the default
         StanzaBase.reply behavior.
+
+        Arguments:
+            clear -- Indicates if existing content should be
+                     removed before replying. Defaults to True.
         """
         self['type'] = 'result'
-        StanzaBase.reply(self)
+        StanzaBase.reply(self, clear)
         return self
 
     def send(self, block=True, timeout=None, callback=None):
diff --git a/sleekxmpp/stanza/message.py b/sleekxmpp/stanza/message.py
index 66c74d8a..6f0cf212 100644
--- a/sleekxmpp/stanza/message.py
+++ b/sleekxmpp/stanza/message.py
@@ -104,7 +104,7 @@ class Message(RootStanza):
         self['type'] = 'normal'
         return self
 
-    def reply(self, body=None):
+    def reply(self, body=None, clear=True):
         """
         Create a message reply.
 
@@ -114,7 +114,9 @@ class Message(RootStanza):
         adds a message body if one is given.
 
         Arguments:
-            body -- Optional text content for the message.
+            body  -- Optional text content for the message.
+            clear -- Indicates if existing content should be removed
+                     before replying. Defaults to True.
         """
         StanzaBase.reply(self)
         if self['type'] == 'groupchat':
diff --git a/sleekxmpp/stanza/presence.py b/sleekxmpp/stanza/presence.py
index 7dcd8f90..60dddf64 100644
--- a/sleekxmpp/stanza/presence.py
+++ b/sleekxmpp/stanza/presence.py
@@ -173,14 +173,18 @@ class Presence(RootStanza):
             # The priority is not a number: we consider it 0 as a default
             return 0
 
-    def reply(self):
+    def reply(self, clear=True):
         """
         Set the appropriate presence reply type.
 
         Overrides StanzaBase.reply.
+
+        Arguments:
+            clear -- Indicates if the stanza contents should be removed
+                     before replying. Defaults to True.
         """
         if self['type'] == 'unsubscribe':
             self['type'] = 'unsubscribed'
         elif self['type'] == 'subscribe':
             self['type'] = 'subscribed'
-        return StanzaBase.reply(self)
+        return StanzaBase.reply(self, clear)
diff --git a/sleekxmpp/stanza/rootstanza.py b/sleekxmpp/stanza/rootstanza.py
index 8123c5f8..bc11476e 100644
--- a/sleekxmpp/stanza/rootstanza.py
+++ b/sleekxmpp/stanza/rootstanza.py
@@ -43,8 +43,8 @@ class RootStanza(StanzaBase):
         Arguments:
             e -- Exception object
         """
-        self.reply()
         if isinstance(e, XMPPError):
+            self.reply(clear=e.clear)
             # We raised this deliberately
             self['error']['condition'] = e.condition
             self['error']['text'] = e.text
@@ -56,6 +56,7 @@ class RootStanza(StanzaBase):
                 self['error']['type'] = e.etype
             self.send()
         else:
+            self.reply()
             # We probably didn't raise this on purpose, so send an error stanza
             self['error']['condition'] = 'undefined-condition'
             self['error']['text'] = "SleekXMPP got into trouble."
diff --git a/sleekxmpp/xmlstream/stanzabase.py b/sleekxmpp/xmlstream/stanzabase.py
index 3937a7a9..1f229cea 100644
--- a/sleekxmpp/xmlstream/stanzabase.py
+++ b/sleekxmpp/xmlstream/stanzabase.py
@@ -1161,12 +1161,17 @@ class StanzaBase(ElementBase):
         self.clear()
         return self
 
-    def reply(self):
+    def reply(self, clear=True):
         """
-        Reset the stanza and swap its 'from' and 'to' attributes to prepare
-        for sending a reply stanza.
+        Swap the 'from' and 'to' attributes to prepare the stanza for
+        sending a reply. If clear=True, then also remove the stanza's
+        contents to make room for the reply content.
 
         For client streams, the 'from' attribute is removed.
+
+        Arguments:
+            clear -- Indicates if the stanza's contents should be
+                     removed. Defaults to True
         """
         # if it's a component, use from
         if self.stream and hasattr(self.stream, "is_component") and \
@@ -1175,7 +1180,8 @@ class StanzaBase(ElementBase):
         else:
             self['to'] = self['from']
             del self['from']
-        self.clear()
+        if clear:
+            self.clear()
         return self
 
     def error(self):
diff --git a/tests/test_stream_exceptions.py b/tests/test_stream_exceptions.py
index e1b70d39..a4598a10 100644
--- a/tests/test_stream_exceptions.py
+++ b/tests/test_stream_exceptions.py
@@ -1,5 +1,7 @@
 import sys
 import sleekxmpp
+from sleekxmpp.xmlstream.matcher import MatchXPath
+from sleekxmpp.xmlstream.handler import Callback
 from sleekxmpp.exceptions import XMPPError
 from sleekxmpp.test import *
 
@@ -46,6 +48,41 @@ class TestStreamExceptions(SleekTest):
           </message>
         """, use_values=False)
 
+    def testIqErrorException(self):
+        """Test using error exceptions with Iq stanzas."""
+
+        def handle_iq(iq):
+            raise XMPPError(condition='feature-not-implemented',
+                            text="We don't do things that way here.",
+                            etype='cancel',
+                            clear=False)
+
+        self.stream_start()
+        self.xmpp.register_handler(
+                Callback(
+                    'Test Iq',
+                     MatchXPath('{%s}iq/{test}query' % self.xmpp.default_ns),
+                     handle_iq))
+
+        self.recv("""
+          <iq type="get" id="0">
+            <query xmlns="test" />
+          </iq>
+        """)
+
+        self.send("""
+          <iq type="error" id="0">
+            <query xmlns="test" />
+            <error type="cancel">
+              <feature-not-implemented
+                  xmlns="urn:ietf:params:xml:ns:xmpp-stanzas" />
+              <text xmlns="urn:ietf:params:xml:ns:xmpp-stanzas">
+                We don&apos;t do things that way here.
+              </text>
+            </error>
+          </iq>
+        """, use_values=False)
+
     def testThreadedXMPPErrorException(self):
         """Test raising an XMPPError exception in a threaded handler."""