use io

Author: alingse

jsoncsv/jsontool.py

@@ -135,6 +135,10 @@ def restore(expobj, separator='.', safe=False):
 
 
 def convert_json(fin, fout, func, separator=".", safe=False, json_array=False):
+    '''
+    ensure fin/fout is TextIO
+    '''
+
     if func not in [expand, restore]:
         raise ValueError("unknow convert_json type")
 
@@ -159,9 +163,5 @@ def gen_objs_from_array():
     for obj in objs:
         new = func(obj, separator=separator, safe=safe)
         content = json.dumps(new, ensure_ascii=False)
-        if PY2:
-            fout.write(content.encode('utf-8'))
-            fout.write('\n'.encode('utf-8'))
-        else:
-            fout.write(content.encode('utf-8'))
-            fout.write(str('\n').encode('utf-8'))
+        fout.write(content)
+        fout.write('\n')

jsoncsv/main.py

@@ -46,21 +46,22 @@ def separator_type(sep):
     '--expand',
     'expand',
     is_flag=True,
-    help='expand json')
+    help='expand json (default True)')
 @click.argument(
     'input',
     type=click.File('r', encoding='utf-8'),
     default='-')
 @click.argument(
     'output',
-    type=click.File('wb'),
+    type=click.File('w', encoding='utf-8'),
     default='-')
 def jsoncsv(output, input, expand, restore, safe, separator, json_array):
     if expand and restore:
         raise click.UsageError('can not choose both, default is `-e`')
 
-    func = jsontool.expand
-    if restore:
+    if not restore:
+        func = jsontool.expand
+    else:
         func = jsontool.restore
 
     convert_json(input, output, func, separator=separator, safe=safe, json_array=json_array)

tests/test_jsontool.py

@@ -104,45 +104,45 @@ def test_expand_and_restore(self):
 class TestConvertJSON(unittest.TestCase):
 
     def test_convert_expand(self):
-        fin = io.BytesIO(u'{"a":{"b":3}}\n{"a":{"c":4}}\n'.encode('utf-8'))
-        fout = io.BytesIO()
+        fin = io.StringIO(u'{"a":{"b":3}}\n{"a":{"c":4}}\n')
+        fout = io.StringIO()
 
         convert_json(fin, fout, expand)
 
-        self.assertEqual(b'{"a.b": 3}\n{"a.c": 4}\n', fout.getvalue())
+        self.assertEqual('{"a.b": 3}\n{"a.c": 4}\n', fout.getvalue())
 
         fin.close()
         fout.close()
 
     def test_convert_with_unicode(self):
-        fin = io.BytesIO(u'{"河流":{"长度":3}}\n{"河流":{"名字":"长江"}}\n'.encode('utf-8'))
-        fout = io.BytesIO()
+        fin = io.StringIO(u'{"河流":{"长度":3}}\n{"河流":{"名字":"长江"}}\n')
+        fout = io.StringIO()
 
         convert_json(fin, fout, expand)
 
-        self.assertEqual(u'{"河流.长度": 3}\n{"河流.名字": "长江"}\n'.encode('utf-8'), fout.getvalue())
+        self.assertEqual(u'{"河流.长度": 3}\n{"河流.名字": "长江"}\n', fout.getvalue())
 
         fin.close()
         fout.close()
 
     def test_convert_restore(self):
-        fin = io.BytesIO(u'{"a.b": 3}\n{"a.c": 4}\n'.encode('utf-8'))
-        fout = io.BytesIO()
+        fin = io.StringIO(u'{"a.b": 3}\n{"a.c": 4}\n')
+        fout = io.StringIO()
 
         convert_json(fin, fout, restore)
 
-        self.assertEqual(b'{"a": {"b": 3}}\n{"a": {"c": 4}}\n', fout.getvalue())
+        self.assertEqual('{"a": {"b": 3}}\n{"a": {"c": 4}}\n', fout.getvalue())
 
         fin.close()
         fout.close()
 
     def test_convert_expand_json_array(self):
-        fin = io.BytesIO(u'[{"a":{"b":3}},{"a":{"c":4}}]'.encode('utf-8'))
-        fout = io.BytesIO()
+        fin = io.StringIO(u'[{"a":{"b":3}},{"a":{"c":4}}]')
+        fout = io.StringIO()
 
         convert_json(fin, fout, expand, json_array=True)
 
-        self.assertEqual(b'{"a.b": 3}\n{"a.c": 4}\n', fout.getvalue())
+        self.assertEqual('{"a.b": 3}\n{"a.c": 4}\n', fout.getvalue())
 
         fin.close()
         fout.close()