XMLStream: factorize the highlight function so it can be used in tests as well

4 files changed, 33 insertions, 34 deletions

diff --git a/slixmpp/test/slixtest.py b/slixmpp/test/slixtest.py
index 36d880a1..2173c87c 100644
--- a/slixmpp/test/slixtest.py
+++ b/slixmpp/test/slixtest.py
@@ -15,7 +15,7 @@ from slixmpp import ClientXMPP, ComponentXMPP
 from slixmpp.stanza import Message, Iq, Presence
 from slixmpp.xmlstream import ET
 from slixmpp.xmlstream import ElementBase
-from slixmpp.xmlstream.tostring import tostring
+from slixmpp.xmlstream.tostring import tostring, highlight
 from slixmpp.xmlstream.matcher import StanzaPath, MatcherId, MatchIDSender
 from slixmpp.xmlstream.matcher import MatchXMLMask, MatchXPath
 
@@ -268,16 +268,16 @@ class SlixTest(unittest.TestCase):
                 stanza3.values = values
 
                 debug = "Three methods for creating stanzas do not match.\n"
-                debug += "Given XML:\n%s\n" % tostring(xml)
-                debug += "Given stanza:\n%s\n" % tostring(stanza.xml)
-                debug += "Generated stanza:\n%s\n" % tostring(stanza2.xml)
-                debug += "Second generated stanza:\n%s\n" % tostring(stanza3.xml)
+                debug += "Given XML:\n%s\n" % highlight(tostring(xml))
+                debug += "Given stanza:\n%s\n" % highlight(tostring(stanza.xml))
+                debug += "Generated stanza:\n%s\n" % highlight(tostring(stanza2.xml))
+                debug += "Second generated stanza:\n%s\n" % highlight(tostring(stanza3.xml))
                 result = self.compare(xml, stanza.xml, stanza2.xml, stanza3.xml)
             else:
                 debug = "Two methods for creating stanzas do not match.\n"
-                debug += "Given XML:\n%s\n" % tostring(xml)
-                debug += "Given stanza:\n%s\n" % tostring(stanza.xml)
-                debug += "Generated stanza:\n%s\n" % tostring(stanza2.xml)
+                debug += "Given XML:\n%s\n" % highlight(tostring(xml))
+                debug += "Given stanza:\n%s\n" % highlight(tostring(stanza.xml))
+                debug += "Generated stanza:\n%s\n" % highlight(tostring(stanza2.xml))
                 result = self.compare(xml, stanza.xml, stanza2.xml)
 
             self.failUnless(result, debug)
@@ -560,13 +560,13 @@ class SlixTest(unittest.TestCase):
         if method == 'exact':
             self.failUnless(self.compare(xml, sent_xml),
                 "Features do not match.\nDesired:\n%s\nReceived:\n%s" % (
-                    tostring(xml), tostring(sent_xml)))
+                    highlight(tostring(xml)), highlight(tostring(sent_xml))))
         elif method == 'mask':
             matcher = MatchXMLMask(xml)
             self.failUnless(matcher.match(sent_xml),
                 "Stanza did not match using %s method:\n" % method + \
-                "Criteria:\n%s\n" % tostring(xml) + \
-                "Stanza:\n%s" % tostring(sent_xml))
+                "Criteria:\n%s\n" % highlight(tostring(xml)) + \
+                "Stanza:\n%s" % highlight(tostring(sent_xml)))
         else:
             raise ValueError("Uknown matching method: %s" % method)
 
diff --git a/slixmpp/xmlstream/__init__.py b/slixmpp/xmlstream/__init__.py
index b9ab6dea..b5302292 100644
--- a/slixmpp/xmlstream/__init__.py
+++ b/slixmpp/xmlstream/__init__.py
@@ -9,9 +9,9 @@
 from slixmpp.jid import JID
 from slixmpp.xmlstream.stanzabase import StanzaBase, ElementBase, ET
 from slixmpp.xmlstream.stanzabase import register_stanza_plugin
-from slixmpp.xmlstream.tostring import tostring
+from slixmpp.xmlstream.tostring import tostring, highlight
 from slixmpp.xmlstream.xmlstream import XMLStream, RESPONSE_TIMEOUT
 
 __all__ = ['JID', 'StanzaBase', 'ElementBase',
-           'ET', 'StateMachine', 'tostring', 'XMLStream',
+           'ET', 'StateMachine', 'tostring', 'highlight', 'XMLStream',
            'RESPONSE_TIMEOUT']
diff --git a/slixmpp/xmlstream/tostring.py b/slixmpp/xmlstream/tostring.py
index 2d99e323..521ee460 100644
--- a/slixmpp/xmlstream/tostring.py
+++ b/slixmpp/xmlstream/tostring.py
@@ -158,3 +158,19 @@ def escape(text, use_cdata=False):
             escaped = map(lambda x : "<![CDATA[%s]]>" % x, text.split("]]>"))
             return "<![CDATA[]]]><![CDATA[]>]]>".join(escaped)
         return text
+
+
+def _get_highlight():
+    try:
+        from pygments import highlight
+        from pygments.lexers import get_lexer_by_name
+        from pygments.formatters import Terminal256Formatter
+
+        LEXER = get_lexer_by_name('xml')
+        FORMATTER = Terminal256Formatter()
+
+        return lambda x: highlight(x, LEXER, FORMATTER)
+    except ImportError:
+        return lambda x: x
+
+highlight = _get_highlight()
diff --git a/slixmpp/xmlstream/xmlstream.py b/slixmpp/xmlstream/xmlstream.py
index 2b165060..820f5586 100644
--- a/slixmpp/xmlstream/xmlstream.py
+++ b/slixmpp/xmlstream/xmlstream.py
@@ -22,27 +22,10 @@ import uuid
 import xml.etree.ElementTree
 
 from slixmpp.xmlstream.asyncio import asyncio
-from slixmpp.xmlstream import tostring
+from slixmpp.xmlstream import tostring, highlight
 from slixmpp.xmlstream.stanzabase import StanzaBase, ElementBase
 from slixmpp.xmlstream.resolver import resolve, default_resolver
 
-try:
-    from pygments import highlight
-    from pygments.lexers import get_lexer_by_name
-    from pygments.formatters import Terminal256Formatter
-
-    LEXER = get_lexer_by_name('xml')
-    FORMATTER = Terminal256Formatter()
-except ImportError:
-    def highlight(text, lexer, formatter, outfile=None):
-        if outfile is not None:
-            outfile.write(text)
-            return
-        return text
-
-    LEXER = None
-    FORMATTER = None
-
 #: The time in seconds to wait before timing out waiting for response stanzas.
 RESPONSE_TIMEOUT = 30
 
@@ -363,7 +346,7 @@ class XMLStream(asyncio.BaseProtocol):
                     log.debug('RECV: %s', highlight(tostring(self.xml_root, xmlns=self.default_ns,
                                                          stream=self,
                                                          top_level=True,
-                                                         open_only=True), LEXER, FORMATTER).strip())
+                                                         open_only=True)).strip())
                     self.start_stream_handler(self.xml_root)
                 self.xml_depth += 1
             if event == 'end':
@@ -853,7 +836,7 @@ class XMLStream(asyncio.BaseProtocol):
 
         :param string data: Any bytes or utf-8 string value.
         """
-        log.debug("SEND: %s", highlight(data, LEXER, FORMATTER).strip())
+        log.debug("SEND: %s", highlight(data).strip())
         if not self.transport:
             raise NotConnectedError()
         if isinstance(data, str):
@@ -906,7 +889,7 @@ class XMLStream(asyncio.BaseProtocol):
         if stanza is None:
             return
 
-        log.debug("RECV: %s", highlight(str(stanza), LEXER, FORMATTER).strip())
+        log.debug("RECV: %s", highlight(str(stanza)).strip())
 
         # Match the stanza against registered handlers. Handlers marked
         # to run "in stream" will be executed immediately; the rest will