support IO with bytes for easier usage on python3

src/ERP5Diff/ERP5Diff.py

@@ -32,6 +32,7 @@ parser = etree.XMLParser(remove_blank_text=True)
 import sys
 import getopt
 import os
+from io import BytesIO
 from six import StringIO
 import re
 import codecs
@@ -97,7 +98,7 @@ class ERP5Diff:
     """
     doc_list = []
     for a in args:
-      if isinstance(a, str):
+      if isinstance(a, six.string_types + (bytes, )):
         doc_list.append(etree.fromstring(a, parser))
       else:
         element_tree = etree.parse(a, parser)
@@ -672,7 +673,7 @@ class ERP5Diff:
       del old_doc
       del new_doc
 
-  def output(self, output_file=None):
+  def output(self, output_file=None, encoding='unicode'):
     """
       Output the result of parsing XML documents to 'output_file'.
       If it is not specified, stdout is assumed.
@@ -680,9 +681,19 @@ class ERP5Diff:
     if output_file is None:
       output_file = sys.stdout
     # stream
-    xml = etree.tostring(self._result, encoding='utf-8', pretty_print=True)
+    xml = etree.tostring(self._result, encoding=encoding, pretty_print=True)
     output_file.write(xml)
 
+  def outputBytes(self, encoding='utf-8'):
+    """
+      Return the result as a bytes object.
+    """
+    io = BytesIO()
+    self.output(io, encoding=encoding)
+    ret = io.getvalue()
+    io.close()
+    return ret
+
   def outputString(self):
     """
       Return the result as a string object.

src/tests/test_erp5diff.py

 # -*- coding: utf-8 -*-
 from __future__ import absolute_import
 import doctest
+import io
 import unittest
 import lxml.doctestcompare
 from lxml import etree
@@ -1153,5 +1154,65 @@ class TestERP5Diff(unittest.TestCase):
 """
     self._assertERP5DiffWorks(old_xml, new_xml, expected_result_string)
 
+class TestIO(unittest.TestCase):
+  def setUp(self):
+    self.erp5diff = ERP5Diff()
+
+  def assertXupdate(self):
+    expected_result_string = """
+    <xupdate:modifications xmlns:xupdate="http://www.xmldb.org/xupdate" version="1.0">
+      <xupdate:update select="/xml/attribute::case">é</xupdate:update>
+    </xupdate:modifications>
+    """
+
+    result_string = self.erp5diff.outputString()
+    checker = lxml.doctestcompare.LXMLOutputChecker()
+    if not checker.check_output(expected_result_string, result_string, 0):
+      self.fail(
+        checker.output_difference(
+          doctest.Example("", expected_result_string),
+          result_string,
+          0))
+
+    self.assertIn(u'é'.encode('utf-8'), self.erp5diff.outputBytes(encoding='utf-8'))
+    out = io.BytesIO()
+    self.erp5diff.output(out, encoding='utf-8')
+    self.assertIn(u'é'.encode('utf-8'), out.getvalue())
+
+    self.assertIn(u'é'.encode('iso-8859-1'), self.erp5diff.outputBytes(encoding='iso-8859-1'))
+    out = io.BytesIO()
+    self.erp5diff.output(out, encoding='iso-8859-1')
+    self.assertIn(u'é'.encode('iso-8859-1'), out.getvalue())
+
+  def test_unicode(self):
+    self.erp5diff.compare(
+      u'<xml case="à"></xml>',
+      u'<xml case="é"></xml>')
+    self.assertXupdate()
+
+  def test_stringio(self):
+    self.erp5diff.compare(
+      io.StringIO(u'<xml case="à"></xml>'),
+      io.StringIO(u'<xml case="é"></xml>'))
+    self.assertXupdate()
+
+  def test_utf8_encoded_bytes(self):
+    self.erp5diff.compare(
+      u'<xml case="à"></xml>'.encode('utf-8'),
+      u'<xml case="é"></xml>'.encode('utf-8'))
+    self.assertXupdate()
+
+  def test_iso88591_encoded_bytes(self):
+    self.erp5diff.compare(
+      u'<?xml version="1.0" encoding="iso-8859-1"?><xml case="à"></xml>'.encode('iso-8859-1'),
+      u'<?xml version="1.0" encoding="iso-8859-1"?><xml case="é"></xml>'.encode('iso-8859-1'))
+    self.assertXupdate()
+
+  def test_utf8_encoded_byteesio(self):
+    self.erp5diff.compare(
+      io.BytesIO(u'<xml case="à"></xml>'.encode('utf-8')),
+      io.BytesIO(u'<xml case="é"></xml>'.encode('utf-8')))
+    self.assertXupdate()
+
 if __name__ == '__main__':
   unittest.main()