test_callable.py 814 B

123456789101112131415161718192021222324252627282930313233
  1. import pandas as pd
  2. import pandas.util.testing as tm
  3. def test_getitem_callable():
  4. # GH 12533
  5. s = pd.Series(4, index=list('ABCD'))
  6. result = s[lambda x: 'A']
  7. assert result == s.loc['A']
  8. result = s[lambda x: ['A', 'B']]
  9. tm.assert_series_equal(result, s.loc[['A', 'B']])
  10. result = s[lambda x: [True, False, True, True]]
  11. tm.assert_series_equal(result, s.iloc[[0, 2, 3]])
  12. def test_setitem_callable():
  13. # GH 12533
  14. s = pd.Series([1, 2, 3, 4], index=list('ABCD'))
  15. s[lambda x: 'A'] = -1
  16. tm.assert_series_equal(s, pd.Series([-1, 2, 3, 4], index=list('ABCD')))
  17. def test_setitem_other_callable():
  18. # GH 13299
  19. inc = lambda x: x + 1
  20. s = pd.Series([1, 2, -1, 4])
  21. s[s < 0] = inc
  22. expected = pd.Series([1, 2, inc, 4])
  23. tm.assert_series_equal(s, expected)