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 : "" % x, text.split("]]>")) return "]]>".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