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

3 files changed, 22 insertions, 23 deletions

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