diff options
| -rw-r--r-- | tablib/formats/_csv.py | 18 | ||||
| -rwxr-xr-x | test_tablib.py | 8 |
2 files changed, 22 insertions, 4 deletions
diff --git a/tablib/formats/_csv.py b/tablib/formats/_csv.py index ddf477a..bfe8b0f 100644 --- a/tablib/formats/_csv.py +++ b/tablib/formats/_csv.py @@ -5,12 +5,17 @@ import sys if sys.version_info[0] > 2: + is_py3 = True + from io import StringIO + import csv else: + is_py3 = False from cStringIO import StringIO + import tablib.packages.unicodecsv as csv + -import tablib.packages.unicodecsv as csv import os import tablib @@ -27,7 +32,11 @@ DEFAULT_ENCODING = 'utf-8' def export_set(dataset): """Returns CSV representation of Dataset.""" stream = StringIO() - _csv = csv.writer(stream, encoding=DEFAULT_ENCODING) + + if is_py3: + _csv = csv.writer(stream) + else: + _csv = csv.writer(stream, encoding=DEFAULT_ENCODING) for row in dataset._package(dicts=False): _csv.writerow(row) @@ -40,7 +49,10 @@ def import_set(dset, in_stream, headers=True): dset.wipe() - rows = csv.reader(in_stream.splitlines(), encoding=DEFAULT_ENCODING) + if is_py3: + rows = csv.reader(in_stream.splitlines()) + else: + rows = csv.reader(in_stream.splitlines(), encoding=DEFAULT_ENCODING) for i, row in enumerate(rows): if (i == 0) and (headers): diff --git a/test_tablib.py b/test_tablib.py index 211c6a7..c7f4754 100755 --- a/test_tablib.py +++ b/test_tablib.py @@ -503,7 +503,13 @@ class TablibTestCase(unittest.TestCase): """Check if unicode in csv export doesn't raise.""" data = tablib.Dataset() - data.append([u'\xfc', u'\xfd']) + + if sys.version_info[0] > 2: + data.append(['\xfc', '\xfd']) + else: + exec("data.append([u'\xfc', u'\xfd'])") + + data.csv |