test_head_tail.py 1.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657
  1. import numpy as np
  2. from pandas import DataFrame
  3. import pandas._testing as tm
  4. def test_head_tail_generic(index, frame_or_series):
  5. # GH#5370
  6. ndim = 2 if frame_or_series is DataFrame else 1
  7. shape = (len(index),) * ndim
  8. vals = np.random.randn(*shape)
  9. obj = frame_or_series(vals, index=index)
  10. tm.assert_equal(obj.head(), obj.iloc[:5])
  11. tm.assert_equal(obj.tail(), obj.iloc[-5:])
  12. # 0-len
  13. tm.assert_equal(obj.head(0), obj.iloc[0:0])
  14. tm.assert_equal(obj.tail(0), obj.iloc[0:0])
  15. # bounded
  16. tm.assert_equal(obj.head(len(obj) + 1), obj)
  17. tm.assert_equal(obj.tail(len(obj) + 1), obj)
  18. # neg index
  19. tm.assert_equal(obj.head(-3), obj.head(len(index) - 3))
  20. tm.assert_equal(obj.tail(-3), obj.tail(len(index) - 3))
  21. def test_head_tail(float_frame):
  22. tm.assert_frame_equal(float_frame.head(), float_frame[:5])
  23. tm.assert_frame_equal(float_frame.tail(), float_frame[-5:])
  24. tm.assert_frame_equal(float_frame.head(0), float_frame[0:0])
  25. tm.assert_frame_equal(float_frame.tail(0), float_frame[0:0])
  26. tm.assert_frame_equal(float_frame.head(-1), float_frame[:-1])
  27. tm.assert_frame_equal(float_frame.tail(-1), float_frame[1:])
  28. tm.assert_frame_equal(float_frame.head(1), float_frame[:1])
  29. tm.assert_frame_equal(float_frame.tail(1), float_frame[-1:])
  30. # with a float index
  31. df = float_frame.copy()
  32. df.index = np.arange(len(float_frame)) + 0.1
  33. tm.assert_frame_equal(df.head(), df.iloc[:5])
  34. tm.assert_frame_equal(df.tail(), df.iloc[-5:])
  35. tm.assert_frame_equal(df.head(0), df[0:0])
  36. tm.assert_frame_equal(df.tail(0), df[0:0])
  37. tm.assert_frame_equal(df.head(-1), df.iloc[:-1])
  38. tm.assert_frame_equal(df.tail(-1), df.iloc[1:])
  39. def test_head_tail_empty():
  40. # test empty dataframe
  41. empty_df = DataFrame()
  42. tm.assert_frame_equal(empty_df.tail(), empty_df)
  43. tm.assert_frame_equal(empty_df.head(), empty_df)