Updated RootStanza to use registerStanzaPlugin, and be PEP8 compliant.

2 files changed, 57 insertions, 26 deletions

diff --git a/sleekxmpp/stanza/__init__.py b/sleekxmpp/stanza/__init__.py
index b8fca228..8302c43d 100644
--- a/sleekxmpp/stanza/__init__.py
+++ b/sleekxmpp/stanza/__init__.py
@@ -5,4 +5,9 @@
 
     See the file LICENSE for copying permission.
 """
-__all__ = ['presence']
+
+
+from sleekxmpp.stanza.error import Error
+from sleekxmpp.stanza.iq import Iq
+from sleekxmpp.stanza.message import Message
+from sleekxmpp.stanza.presence import Presence
diff --git a/sleekxmpp/stanza/rootstanza.py b/sleekxmpp/stanza/rootstanza.py
index 3bae1f0d..fb5498f5 100644
--- a/sleekxmpp/stanza/rootstanza.py
+++ b/sleekxmpp/stanza/rootstanza.py
@@ -5,34 +5,60 @@
 
     See the file LICENSE for copying permission.
 """
-from .. xmlstream.stanzabase import StanzaBase
-from xml.etree import cElementTree as ET
-from . error import Error
-from .. exceptions import XMPPError
+
 import logging
 import traceback
 import sys
 
+from sleekxmpp.exceptions import XMPPError
+from sleekxmpp.stanza import Error
+from sleekxmpp.xmlstream.stanzabase import ET, StanzaBase, registerStanzaPlugin
+
+
 class RootStanza(StanzaBase):
 
-	def exception(self, e): #called when a handler raises an exception
-		self.reply()
-		if isinstance(e, XMPPError): # we raised this deliberately
-			self['error']['condition'] = e.condition
-			self['error']['text'] = e.text
-			if e.extension is not None: # extended error tag
-				extxml = ET.Element("{%s}%s" % (e.extension_ns, e.extension), e.extension_args)
-				self['error'].xml.append(extxml)
-				self['error']['type'] = e.etype
-		else: # we probably didn't raise this on purpose, so send back a traceback
-			self['error']['condition'] = 'undefined-condition'
-			if sys.version_info < (3,0):
-				self['error']['text'] = "SleekXMPP got into trouble."
-			else:
-				self['error']['text'] = traceback.format_tb(e.__traceback__)
-			logging.exception('Error handling {%s}%s stanza' % (self.namespace, self.name))	
-		self.send()
-
-# all jabber:client root stanzas should have the error plugin
-RootStanza.plugin_attrib_map['error'] = Error
-RootStanza.plugin_tag_map["{%s}%s" % (Error.namespace, Error.name)] = Error
+    """
+    A top-level XMPP stanza in an XMLStream.
+
+    The RootStanza class provides a more XMPP specific exception
+    handler than provided by the generic StanzaBase class.
+
+    Methods:
+        exception -- Overrides StanzaBase.exception
+    """
+
+    def exception(self, e):
+        """
+        Create and send an error reply.
+
+        Typically called when an event handler raises an exception.
+        The error's type and text content are based on the exception
+        object's type and content.
+
+        Overrides StanzaBase.exception.
+
+        Arguments:
+            e -- Exception object
+        """
+        self.reply()
+        if isinstance(e, XMPPError): 
+            # We raised this deliberately
+            self['error']['condition'] = e.condition
+            self['error']['text'] = e.text
+            if e.extension is not None: 
+                # Extended error tag
+                extxml = ET.Element("{%s}%s" % (e.extension_ns, e.extension), e.extension_args)
+                self['error'].append(extxml)
+                self['error']['type'] = e.etype
+        else: 
+            # We probably didn't raise this on purpose, so send back a traceback
+            self['error']['condition'] = 'undefined-condition'
+            if sys.version_info < (3,0):
+                self['error']['text'] = "SleekXMPP got into trouble."
+            else:
+                self['error']['text'] = traceback.format_tb(e.__traceback__)
+                logging.exception('Error handling {%s}%s stanza' % (self.namespace, self.name))	
+        self.send()
+
+
+registerStanzaPlugin(RootStanza, Error)