parametrize tests and allow index parameter to handle MultiIndex and index name

Author: reidy-p

pandas/core/generic.py

@@ -1673,9 +1673,9 @@ def to_json(self, path_or_buf=None, orient=None, date_format=None,
             .. versionadded:: 0.21.0
 
         index : boolean, default True
-            Whether to include the index values in the JSON string. A
-            ValueError will be thrown if index is False when orient is not
-            'split' or 'table'.
+            Whether to include the index values in the JSON string. Not
+            including the index (``index=False``) is only supported when
+            orient is 'split' or 'table'.
 
             .. versionadded:: 0.22.0
 

pandas/io/json/json.py

@@ -200,15 +200,17 @@ def __init__(self, obj, orient, date_format, double_precision,
         if is_period_dtype(obj.index):
             obj.index = obj.index.to_timestamp()
 
-        self.obj = obj.reset_index()
+        # exclude index from obj if index=False
+        if not self.index:
+            self.obj = obj.reset_index(drop=True)
+        else:
+            self.obj = obj.reset_index(drop=False)
         self.date_format = 'iso'
         self.orient = 'records'
         self.index = index
 
     def _write(self, obj, orient, double_precision, ensure_ascii,
                date_unit, iso_dates, default_handler):
-        if not self.index:
-            obj = obj.drop('index', axis=1)
         data = super(JSONTableWriter, self)._write(obj, orient,
                                                    double_precision,
                                                    ensure_ascii, date_unit,

pandas/tests/io/json/test_pandas.py

@@ -1149,80 +1149,60 @@ def test_data_frame_size_after_to_json(self):
 
         assert size_before == size_after
 
-    def test_index_false_to_json(self):
+    @pytest.mark.parametrize('data, expected', [
+        (DataFrame([[1, 2], [4, 5]], columns=['a', 'b']),
+            {'columns': ['a', 'b'], 'data': [[1, 2], [4, 5]]}),
+        (DataFrame([[1, 2], [4, 5]], columns=['a', 'b']).rename_axis('foo'),
+            {'columns': ['a', 'b'], 'data': [[1, 2], [4, 5]]}),
+        (DataFrame([[1, 2], [4, 5]], columns=['a', 'b'],
+                   index=[['a', 'b'], ['c', 'd']]),
+            {'columns': ['a', 'b'], 'data': [[1, 2], [4, 5]]}),
+        (Series([1, 2, 3], name='A'),
+            {'name': 'A', 'data': [1, 2, 3]}),
+        (Series([1, 2, 3], name='A').rename_axis('foo'),
+            {'name': 'A', 'data': [1, 2, 3]}),
+        (Series([1, 2], name='A', index=[['a', 'b'], ['c', 'd']]),
+            {'name': 'A', 'data': [1, 2]}),
+    ])
+    def test_index_false_to_json_split(self, data, expected):
         # GH 17394
-        # Testing index parameter in to_json
-
-        # Testing DataFrame.to_json(orient='split', index=False)
-        df = pd.DataFrame([[1, 2], [4, 5]], columns=['a', 'b'])
-
-        result = df.to_json(orient='split', index=False)
-        result = json.loads(result)
-
-        expected = {
-            'columns': ['a', 'b'],
-            'data': [[1, 2], [4, 5]]
-        }
-
-        assert result == expected
+        # Testing index=False in to_json with orient='split'
 
-        # Testing DataFrame.to_json(orient='table', index=False)
-        result = df.to_json(orient='table', index=False)
+        result = data.to_json(orient='split', index=False)
         result = json.loads(result)
 
-        schema = {
-            'fields': [{'name': 'a', 'type': 'integer'},
-                       {'name': 'b', 'type': 'integer'}],
-            'pandas_version': '0.20.0'
-        }
-
-        expected = {
-            'schema': schema,
-            'data': [{'a': 1, 'b': 2}, {'a': 4, 'b': 5}]
-        }
-
         assert result == expected
 
-        # Testing Series.to_json(orient='split', index=False)
-        s = pd.Series([1, 2, 3], name='A')
-
-        result = s.to_json(orient='split', index=False)
-        result = json.loads(result)
-
-        expected = {
-            'name': 'A',
-            'data': [1, 2, 3]
-        }
-
-        assert result == expected
+    @pytest.mark.parametrize('data', [
+        (DataFrame([[1, 2], [4, 5]], columns=['a', 'b'])),
+        (DataFrame([[1, 2], [4, 5]], columns=['a', 'b']).rename_axis('foo')),
+        (DataFrame([[1, 2], [4, 5]], columns=['a', 'b'],
+                   index=[['a', 'b'], ['c', 'd']])),
+        (Series([1, 2, 3], name='A')),
+        (Series([1, 2, 3], name='A').rename_axis('foo')),
+        (Series([1, 2], name='A', index=[['a', 'b'], ['c', 'd']])),
+    ])
+    def test_index_false_to_json_table(self, data):
+        # GH 17394
+        # Testing index=False in to_json with orient='table'
 
-        # Testing Series.to_json(orient='table', index=False)
-        result = s.to_json(orient='table', index=False)
+        result = data.to_json(orient='table', index=False)
         result = json.loads(result)
 
-        fields = [{'name': 'A', 'type': 'integer'}]
-
-        schema = {
-            'fields': fields,
-            'pandas_version': '0.20.0'
-        }
-
         expected = {
-            'schema': schema,
-            'data': [{'A': 1}, {'A': 2}, {'A': 3}]
+            'schema': pd.io.json.build_table_schema(data, index=False),
+            'data': DataFrame(data).to_dict(orient='records')
         }
 
         assert result == expected
 
     @pytest.mark.parametrize('orient', [
-        ('records'),
-        ('index'),
-        ('columns'),
-        ('values'),
+        'records', 'index', 'columns', 'values'
     ])
     def test_index_false_error_to_json(self, orient):
         # GH 17394
         # Testing error message from to_json with index=False
+
         df = pd.DataFrame([[1, 2], [4, 5]], columns=['a', 'b'])
 
         with tm.assert_raises_regex(ValueError, "'index=False' is only "