Fix escape_string doesn't accept unicode on Python 2

methane · methane · commit 0fda4ac992a1 · 2016-01-14T02:13:21.000+09:00
fixes PyMySQL#415
diff --git a/pymysql/_compat.py b/pymysql/_compat.py
@@ -7,12 +7,15 @@
 CPYTHON = not PYPY and not JYTHON and not IRONPYTHON
 
 if PY2:
+    import __builtin__
     range_type = xrange
     text_type = unicode
     long_type = long
     str_type = basestring
+    unichr = __builtin__.unichr
 else:
     range_type = range
     text_type = str
     long_type = int
     str_type = str
+    unichr = chr
diff --git a/pymysql/converters.py b/pymysql/converters.py
@@ -1,7 +1,5 @@
-from ._compat import PY2, text_type, long_type, JYTHON, IRONPYTHON
+from ._compat import PY2, text_type, long_type, JYTHON, IRONPYTHON, unichr
 
-import sys
-import binascii
 import datetime
 from decimal import Decimal
 import time
@@ -57,7 +55,7 @@ def escape_int(value, mapping=None):
 def escape_float(value, mapping=None):
     return ('%.15g' % value)
 
-_escape_table = [chr(x) for x in range(128)]
+_escape_table = [unichr(x) for x in range(128)]
 _escape_table[0] = u'\\0'
 _escape_table[ord('\\')] = u'\\\\'
 _escape_table[ord('\n')] = u'\\n'
@@ -80,7 +78,7 @@ def escape_string(value, mapping=None):
         Value should be bytes or unicode.
         """
         if isinstance(value, unicode):
-            return escape_unicode(value)
+            return _escape_unicode(value)
         assert isinstance(value, (bytes, bytearray))
         value = value.replace('\\', '\\\\')
         value = value.replace('\0', '\\0')
diff --git a/pymysql/tests/__init__.py b/pymysql/tests/__init__.py
@@ -6,6 +6,7 @@
 from pymysql.tests.test_SSCursor import *
 from pymysql.tests.test_load_local import *
 from pymysql.tests.test_optionfile import *
+from pymysql.tests.test_converters import *
 
 from pymysql.tests.thirdparty import *
 
diff --git a/pymysql/tests/test_converters.py b/pymysql/tests/test_converters.py
@@ -0,0 +1,23 @@
+from unittest import TestCase
+
+from pymysql._compat import PY2
+from pymysql import converters
+
+
+__all__ = ["TestConverter"]
+
+
+class TestConverter(TestCase):
+
+    def test_escape_string(self):
+        self.assertEqual(
+            converters.escape_string(u"foo\nbar"),
+            u"foo\\nbar"
+        )
+
+    if PY2:
+        def test_escape_string_bytes(self):
+            self.assertEqual(
+                converters.escape_string(b"foo\nbar"),
+                b"foo\\nbar"
+            )
