Merge pull request #224 from CybOXProject/unicode

More unicode and encoding support.

Author: gtback

cybox/__init__.py

@@ -8,6 +8,7 @@
 import json
 from StringIO import StringIO
 
+import cybox.bindings as bindings
 import cybox.utils.idgen
 from cybox.utils import Namespace, META
 
@@ -233,8 +234,12 @@ def from_dict(cls, cls_dict=None):
         return entity
 
     def to_xml(self, include_namespaces=True, namespace_dict=None,
-               pretty=True):
-        """Export an object as an XML String.
+               pretty=True, encoding='utf-8'):
+        """Serializes a :class:`Entity` instance to an XML string.
+
+        The default character encoding is ``utf-8`` and can be set via the
+        `encoding` parameter. If `encoding` is ``None``, a unicode string
+        is returned.
 
         Args:
             include_namespaces (bool): whether to include xmlns and
@@ -244,9 +249,13 @@ def to_xml(self, include_namespaces=True, namespace_dict=None,
                 prefixes
             pretty (bool): whether to produce readable (``True``) or compact
                 (``False``) output. Defaults to ``True``.
+            encoding: The output character encoding. Default is ``utf-8``. If
+                `encoding` is set to ``None``, a unicode string is returned.
 
         Returns:
-            XML string
+            An XML string for this
+            :class:`Entity` instance. Default character encoding is ``utf-8``.
+
         """
         namespace_def = ""
 
@@ -256,10 +265,22 @@ def to_xml(self, include_namespaces=True, namespace_dict=None,
         if not pretty:
             namespace_def = namespace_def.replace('\n\t', ' ')
 
-        s = StringIO()
-        self.to_obj().export(s.write, 0, namespacedef_=namespace_def,
-                             pretty_print=pretty)
-        return s.getvalue().strip()
+
+        with bindings.save_encoding(encoding):
+            sio = StringIO()
+            self.to_obj().export(
+                sio.write,
+                0,
+                namespacedef_=namespace_def,
+                pretty_print=pretty
+            )
+
+        s = unicode(sio.getvalue()).strip()
+
+        if encoding:
+            return s.encode(encoding)
+
+        return s
 
     def to_json(self):
         """Export an object as a JSON String."""

cybox/bindings/__init__.py

@@ -4,6 +4,7 @@
 import base64
 from datetime import datetime, tzinfo, timedelta
 import re
+import contextlib
 
 from xml.sax import saxutils
 from lxml import etree as etree_
@@ -19,6 +20,18 @@
 _Tag_strip_pattern_ = re.compile(r'\{.*\}')
 
 
+@contextlib.contextmanager
+def save_encoding(encoding='utf-8'):
+    global ExternalEncoding
+
+    try:
+        orig_encoding = ExternalEncoding
+        ExternalEncoding = encoding
+        yield
+    finally:
+        ExternalEncoding = orig_encoding
+
+
 def parsexml_(*args, **kwargs):
     if 'parser' not in kwargs:
         # Use the lxml ElementTree compatible parser so that, e.g.,
@@ -258,7 +271,7 @@ def showIndent(lwrite, level, pretty_print=True):
 
 def quote_xml(text):
     if text is None:
-        return ''
+        return u''
 
     # Convert `text` to unicode string. This is mainly a catch-all for non
     # string/unicode types like bool and int.
@@ -267,9 +280,6 @@ def quote_xml(text):
     except UnicodeDecodeError:
         text = text.decode(ExternalEncoding)
 
-    # Convert unicode string to correct output character encoding.
-    text = text.encode(ExternalEncoding)
-
     # If it's a CDATA block, return the text as is.
     if text.startswith(CDATA_START):
         return text
@@ -281,7 +291,7 @@ def quote_xml(text):
 
 def quote_attrib(text):
     if text is None:
-        return '""'  # Return an empty XML attribute value
+        return u'""'
 
     # Convert `text` to unicode string. This is mainly a catch-all for non
     # string/unicode types like bool and int.
@@ -290,9 +300,6 @@ def quote_attrib(text):
     except UnicodeDecodeError:
         text = text.decode(ExternalEncoding)
 
-    # Convert the unicode string to the correct output character encoding.
-    text = text.encode(ExternalEncoding)
-
     # Return the escaped the value of text.
     # Note: This wraps the escaped text in quotation marks.
     return saxutils.quoteattr(text)

cybox/test/__init__.py

@@ -4,6 +4,7 @@
 import json
 import unittest
 
+import cybox.bindings as bindings
 from cybox import Entity, EntityList, TypedField
 import cybox.bindings.cybox_core as core_binding
 from cybox.core import Observables
@@ -83,7 +84,10 @@ def round_trip(o, output=False, list_=False):
     xobj = o2.to_obj()
 
     # 6. Bindings Object -> XML String
-    xml_string = o2.to_xml(include_namespaces=True)
+    xml_string = o2.to_xml(encoding=bindings.ExternalEncoding)
+
+    if not isinstance(xml_string, unicode):
+        xml_string = xml_string.decode(bindings.ExternalEncoding)
 
     if output:
         print(xml_string)

cybox/test/encoding_test.py

@@ -55,7 +55,6 @@ def test_whois(self):
 
     def test_quote_xml(self):
         s = bindings.quote_xml(UNICODE_STR)
-        s = s.decode(bindings.ExternalEncoding)
         self.assertEqual(s, UNICODE_STR)
 
     def test_quote_attrib(self):
@@ -69,77 +68,89 @@ def test_quote_attrib(self):
         """
         s = bindings.quote_attrib(UNICODE_STR)
         s = s[1:-1]
-        s = s.decode(bindings.ExternalEncoding)
         self.assertEqual(s, UNICODE_STR)
 
     def test_quote_attrib_int(self):
         i = 65536
         s = bindings.quote_attrib(i)
-        s = s[1:-1]
-        self.assertEqual(str(i), s)
+        self.assertEqual(u'"65536"', s)
 
     def test_quote_attrib_bool(self):
         b = True
         s = bindings.quote_attrib(b)
-        s = s[1:-1]
-        self.assertEqual(str(b), s)
+        self.assertEqual(u'"True"', s)
 
     def test_quote_xml_int(self):
         i = 65536
         s = bindings.quote_xml(i)
-        self.assertEqual(str(i), s)
+        self.assertEqual(unicode(i), s)
 
     def test_quote_xml_bool(self):
         b = True
         s = bindings.quote_xml(b)
-        self.assertEqual(str(b), s)
+        self.assertEqual(unicode(b), s)
 
     def test_quote_xml_encoded(self):
         encoding = bindings.ExternalEncoding
         encoded = UNICODE_STR.encode(encoding)
         quoted = bindings.quote_xml(encoded)
-        decoded = quoted.decode(encoding)
-        self.assertEqual(UNICODE_STR, decoded)
+        self.assertEqual(UNICODE_STR, quoted)
 
     def test_quote_attrib_encoded(self):
         encoding = bindings.ExternalEncoding
         encoded = UNICODE_STR.encode(encoding)
         quoted = bindings.quote_attrib(encoded)[1:-1]
-        decoded = quoted.decode(encoding)
-        self.assertEqual(UNICODE_STR, decoded)
+        self.assertEqual(UNICODE_STR, quoted)
 
     def test_quote_xml_zero(self):
         i = 0
         s = bindings.quote_xml(i)
-        self.assertEqual(str(i), s)
+        self.assertEqual(unicode(i), s)
 
     def test_quote_attrib_zero(self):
         i = 0
         s = bindings.quote_attrib(i)
-        s = s[1:-1]
-        self.assertEqual(str(i), s)
+        self.assertEqual(u'"0"', s)
 
     def test_quote_xml_none(self):
         i = None
         s = bindings.quote_xml(i)
-        self.assertEqual('', s)
+        self.assertEqual(u'', s)
 
     def test_quote_attrib_none(self):
         i = None
         s = bindings.quote_attrib(i)
-        s = s[1:-1]
-        self.assertEqual('', s)
+        self.assertEqual(u'""', s)
 
     def test_quote_attrib_empty(self):
         i = ''
         s = bindings.quote_attrib(i)
-        s = s[1:-1]
-        self.assertEqual('', s)
+        self.assertEqual(u'""', s)
 
     def test_quote_xml_empty(self):
         i = ''
         s = bindings.quote_xml(i)
-        self.assertEqual('', s)
+        self.assertEqual(u'', s)
+
+    def test_to_xml_utf16_encoded(self):
+        encoding = 'utf-16'
+        o = Observable()
+        o.title = UNICODE_STR
+        xml = o.to_xml(encoding=encoding)
+        self.assertTrue(UNICODE_STR in xml.decode(encoding))
+
+    def test_to_xml_default_encoded(self):
+        o = Observable()
+        o.title = UNICODE_STR
+        xml = o.to_xml()
+        self.assertTrue(UNICODE_STR in xml.decode('utf-8'))
+
+    def test_to_xml_no_encoding(self):
+        o = Observable()
+        o.title = UNICODE_STR
+        xml = o.to_xml(encoding=None)
+        self.assertTrue(isinstance(xml, unicode))
+        self.assertTrue(UNICODE_STR in xml)
 
 if __name__ == "__main__":
     unittest.main()