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

@@ -1139,90 +1139,70 @@ def roundtrip(s, encoding='latin-1'):
  for s in examples:
  roundtrip(s)
 
- def test_data_frame_size_after_to_json(self):
- # GH15344
- df = DataFrame({'a': [str(1)]})
-
- size_before = df.memory_usage(index=True, deep=True).sum()
- df.to_json()
- size_after = df.memory_usage(index=True, deep=True).sum()
-
- assert size_before == size_after
-
- def test_index_false_to_json(self):
+ def test_data_frame_size_after_to_json(self):
+ # GH15344
+ df = DataFrame({'a': [str(1)]})
+
+ size_before = df.memory_usage(index=True, deep=True).sum()
+ df.to_json()
+ size_after = df.memory_usage(index=True, deep=True).sum()
+
+ assert size_before == size_after
+
+ @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 index=False in to_json with orient='split'
 
- # 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 = data.to_json(orient='split', index=False)
  result = json.loads(result)
 
- expected = {
- 'columns': ['a', 'b'],
- 'data': [[1, 2], [4, 5]]
- }
-
  assert result == expected
 
- # Testing DataFrame.to_json(orient='table', index=False)
- result = df.to_json(orient='table', 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 "