test_subclass.py 4.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108
  1. # coding=utf-8
  2. # pylint: disable-msg=E1101,W0612
  3. import numpy as np
  4. import pandas as pd
  5. from pandas import SparseDtype
  6. import pandas.util.testing as tm
  7. class TestSeriesSubclassing(object):
  8. def test_indexing_sliced(self):
  9. s = tm.SubclassedSeries([1, 2, 3, 4], index=list('abcd'))
  10. res = s.loc[['a', 'b']]
  11. exp = tm.SubclassedSeries([1, 2], index=list('ab'))
  12. tm.assert_series_equal(res, exp)
  13. res = s.iloc[[2, 3]]
  14. exp = tm.SubclassedSeries([3, 4], index=list('cd'))
  15. tm.assert_series_equal(res, exp)
  16. res = s.loc[['a', 'b']]
  17. exp = tm.SubclassedSeries([1, 2], index=list('ab'))
  18. tm.assert_series_equal(res, exp)
  19. def test_to_frame(self):
  20. s = tm.SubclassedSeries([1, 2, 3, 4], index=list('abcd'), name='xxx')
  21. res = s.to_frame()
  22. exp = tm.SubclassedDataFrame({'xxx': [1, 2, 3, 4]}, index=list('abcd'))
  23. tm.assert_frame_equal(res, exp)
  24. def test_subclass_unstack(self):
  25. # GH 15564
  26. s = tm.SubclassedSeries(
  27. [1, 2, 3, 4], index=[list('aabb'), list('xyxy')])
  28. res = s.unstack()
  29. exp = tm.SubclassedDataFrame(
  30. {'x': [1, 3], 'y': [2, 4]}, index=['a', 'b'])
  31. tm.assert_frame_equal(res, exp)
  32. class TestSparseSeriesSubclassing(object):
  33. def test_subclass_sparse_slice(self):
  34. # int64
  35. s = tm.SubclassedSparseSeries([1, 2, 3, 4, 5])
  36. exp = tm.SubclassedSparseSeries([2, 3, 4], index=[1, 2, 3])
  37. tm.assert_sp_series_equal(s.loc[1:3], exp)
  38. assert s.loc[1:3].dtype == SparseDtype(np.int64)
  39. exp = tm.SubclassedSparseSeries([2, 3], index=[1, 2])
  40. tm.assert_sp_series_equal(s.iloc[1:3], exp)
  41. assert s.iloc[1:3].dtype == SparseDtype(np.int64)
  42. exp = tm.SubclassedSparseSeries([2, 3], index=[1, 2])
  43. tm.assert_sp_series_equal(s[1:3], exp)
  44. assert s[1:3].dtype == SparseDtype(np.int64)
  45. # float64
  46. s = tm.SubclassedSparseSeries([1., 2., 3., 4., 5.])
  47. exp = tm.SubclassedSparseSeries([2., 3., 4.], index=[1, 2, 3])
  48. tm.assert_sp_series_equal(s.loc[1:3], exp)
  49. assert s.loc[1:3].dtype == SparseDtype(np.float64)
  50. exp = tm.SubclassedSparseSeries([2., 3.], index=[1, 2])
  51. tm.assert_sp_series_equal(s.iloc[1:3], exp)
  52. assert s.iloc[1:3].dtype == SparseDtype(np.float64)
  53. exp = tm.SubclassedSparseSeries([2., 3.], index=[1, 2])
  54. tm.assert_sp_series_equal(s[1:3], exp)
  55. assert s[1:3].dtype == SparseDtype(np.float64)
  56. def test_subclass_sparse_addition(self):
  57. s1 = tm.SubclassedSparseSeries([1, 3, 5])
  58. s2 = tm.SubclassedSparseSeries([-2, 5, 12])
  59. exp = tm.SubclassedSparseSeries([-1, 8, 17])
  60. tm.assert_sp_series_equal(s1 + s2, exp)
  61. s1 = tm.SubclassedSparseSeries([4.0, 5.0, 6.0])
  62. s2 = tm.SubclassedSparseSeries([1.0, 2.0, 3.0])
  63. exp = tm.SubclassedSparseSeries([5., 7., 9.])
  64. tm.assert_sp_series_equal(s1 + s2, exp)
  65. def test_subclass_sparse_to_frame(self):
  66. s = tm.SubclassedSparseSeries([1, 2], index=list('ab'), name='xxx')
  67. res = s.to_frame()
  68. exp_arr = pd.SparseArray([1, 2], dtype=np.int64, kind='block',
  69. fill_value=0)
  70. exp = tm.SubclassedSparseDataFrame({'xxx': exp_arr},
  71. index=list('ab'),
  72. default_fill_value=0)
  73. tm.assert_sp_frame_equal(res, exp)
  74. # create from int dict
  75. res = tm.SubclassedSparseDataFrame({'xxx': [1, 2]},
  76. index=list('ab'),
  77. default_fill_value=0)
  78. tm.assert_sp_frame_equal(res, exp)
  79. s = tm.SubclassedSparseSeries([1.1, 2.1], index=list('ab'),
  80. name='xxx')
  81. res = s.to_frame()
  82. exp = tm.SubclassedSparseDataFrame({'xxx': [1.1, 2.1]},
  83. index=list('ab'))
  84. tm.assert_sp_frame_equal(res, exp)