Merge pull request #541 from onyxfish/csvjson

csvjson: Use agate.Table.to_json(), closes #539

csvkit/utilities/csvjson.py

@@ -69,11 +69,11 @@ def dump_json(data, newline=False):
         if self.args.streamOutput and (self.args.lat or self.args.lon or self.args.key):
             self.argparser.error('--stream is only allowed if --lat, --lon and --key are not specified.')
 
-        rows = agate.reader(self.input_file, **self.reader_kwargs)
-        column_names = next(rows)
-
         # GeoJSON
         if self.args.lat and self.args.lon:
+            rows = agate.reader(self.input_file, **self.reader_kwargs)
+            column_names = next(rows)
+
             features = []
             min_lon = None
             min_lat = None
@@ -150,40 +150,9 @@ def dump_json(data, newline=False):
                     })
                 ])
             dump_json(output)
-        # Keyed JSON
-        elif self.args.key:
-            output = OrderedDict()
-
-            for row in rows:
-                data = OrderedDict()
-
-                for i, column in enumerate(column_names):
-                    data[column] = row[i]
-
-                k = data[self.args.key]
-
-                if k in output:
-                    raise NonUniqueKeyColumnException('Value %s is not unique in the key column.' % six.text_type(k))
-
-                output[k] = data
-            dump_json(output)
-        # Boring JSON
         else:
-            output = []
-            for row in rows:
-                data = OrderedDict()
-
-                for i, column in enumerate(column_names):
-                    try:
-                        data[column] = row[i]
-                    except IndexError:
-                        data[column] = None
-                if(self.args.streamOutput):
-                    dump_json(data, newline=True)
-                else:
-                    output.append(data)
-            if not self.args.streamOutput:
-                dump_json(output)
+            table = agate.Table.from_csv(self.input_file)
+            table.to_json(stream, key=self.args.key, newline=self.args.streamOutput, indent=self.args.indent)
 
 
 def launch_new_instance():

tests/test_utilities/test_csvjson.py

@@ -31,7 +31,7 @@ def test_simple(self):
 
         js = json.loads(output_file.getvalue())
 
-        self.assertDictEqual(js[0], {"a": "1", "c": "3", "b": "2"})
+        self.assertDictEqual(js[0], {'a': True, 'c': 3.0, 'b': 2.0})
 
     def test_indentation(self):
         args = ['-i', '4', 'examples/dummy.csv']
@@ -40,9 +40,11 @@ def test_indentation(self):
         utility = CSVJSON(args, output_file)
         utility.main()
 
-        js = json.loads(output_file.getvalue())
+        content = output_file.getvalue()
+        js = json.loads(content)
 
-        self.assertDictEqual(js[0], {"a": "1", "c": "3", "b": "2"})
+        self.assertDictEqual(js[0], {'a': True, 'c': 3.0, 'b': 2.0})
+        self.assertRegexpMatches(content, '        "a": true,')
 
     def test_keying(self):
         args = ['-k', 'a', 'examples/dummy.csv']
@@ -53,15 +55,15 @@ def test_keying(self):
 
         js = json.loads(output_file.getvalue())
 
-        self.assertDictEqual(js, {"1": {"a": "1", "c": "3", "b": "2"}})
+        self.assertDictEqual(js, {'true': {'a': True, 'c': 3.0, 'b': 2.0}})
 
     def test_duplicate_keys(self):
         args = ['-k', 'a', 'examples/dummy3.csv']
         output_file = six.StringIO()
 
         utility = CSVJSON(args, output_file)
 
-        self.assertRaises(NonUniqueKeyColumnException, utility.main)
+        self.assertRaisesRegexp(ValueError, 'Value True is not unique in the key column\.', utility.main)
 
     def test_geojson(self):
         args = ['--lat', 'latitude', '--lon', 'longitude', 'examples/test_geo.csv']
@@ -141,5 +143,5 @@ def test_json_streaming(self):
 
         result = list(map(json.loads, output_file.getvalue().splitlines()))
         self.assertEqual(len(result), 2)
-        self.assertDictEqual(result[0], {"a": "1", "c": "3", "b": "2"})
-        self.assertDictEqual(result[1], {"a": "1", "c": "5", "b": "4"})
+        self.assertDictEqual(result[0], {'a': True, 'c': 3.0, 'b': 2.0})
+        self.assertDictEqual(result[1], {'a': True, 'c': 5.0, 'b': 4.0})