/pandas/tests/indexing/test_callable.py

# -*- coding: utf-8 -*-
# pylint: disable-msg=W0612,E1101

import numpy as np
import pandas as pd
import pandas.util.testing as tm


class TestIndexingCallable(object):

    def test_frame_loc_ix_callable(self):
        # GH 11485
        df = pd.DataFrame({'A': [1, 2, 3, 4], 'B': list('aabb'),
                           'C': [1, 2, 3, 4]})
        # iloc cannot use boolean Series (see GH3635)

        # return bool indexer
        res = df.loc[lambda x: x.A > 2]
        tm.assert_frame_equal(res, df.loc[df.A > 2])

        res = df.loc[lambda x: x.A > 2]
        tm.assert_frame_equal(res, df.loc[df.A > 2])

        res = df.loc[lambda x: x.A > 2, ]
        tm.assert_frame_equal(res, df.loc[df.A > 2, ])

        res = df.loc[lambda x: x.A > 2, ]
        tm.assert_frame_equal(res, df.loc[df.A > 2, ])

        res = df.loc[lambda x: x.B == 'b', :]
        tm.assert_frame_equal(res, df.loc[df.B == 'b', :])

        res = df.loc[lambda x: x.B == 'b', :]
        tm.assert_frame_equal(res, df.loc[df.B == 'b', :])

        res = df.loc[lambda x: x.A > 2, lambda x: x.columns == 'B']
        tm.assert_frame_equal(res, df.loc[df.A > 2, [False, True, False]])

        res = df.loc[lambda x: x.A > 2, lambda x: x.columns == 'B']
        tm.assert_frame_equal(res, df.loc[df.A > 2, [False, True, False]])

        res = df.loc[lambda x: x.A > 2, lambda x: 'B']
        tm.assert_series_equal(res, df.loc[df.A > 2, 'B'])

        res = df.loc[lambda x: x.A > 2, lambda x: 'B']
        tm.assert_series_equal(res, df.loc[df.A > 2, 'B'])

        res = df.loc[lambda x: x.A > 2, lambda x: ['A', 'B']]
        tm.assert_frame_equal(res, df.loc[df.A > 2, ['A', 'B']])

        res = df.loc[lambda x: x.A > 2, lambda x: ['A', 'B']]
        tm.assert_frame_equal(res, df.loc[df.A > 2, ['A', 'B']])

        res = df.loc[lambda x: x.A == 2, lambda x: ['A', 'B']]
        tm.assert_frame_equal(res, df.loc[df.A == 2, ['A', 'B']])

        res = df.loc[lambda x: x.A == 2, lambda x: ['A', 'B']]
        tm.assert_frame_equal(res, df.loc[df.A == 2, ['A', 'B']])

        # scalar
        res = df.loc[lambda x: 1, lambda x: 'A']
        assert res == df.loc[1, 'A']

        res = df.loc[lambda x: 1, lambda x: 'A']
        assert res == df.loc[1, 'A']

    def test_frame_loc_ix_callable_mixture(self):
        # GH 11485
        df = pd.DataFrame({'A': [1, 2, 3, 4], 'B': list('aabb'),
                           'C': [1, 2, 3, 4]})

        res = df.loc[lambda x: x.A > 2, ['A', 'B']]
        tm.assert_frame_equal(res, df.loc[df.A > 2, ['A', 'B']])

        res = df.loc[lambda x: x.A > 2, ['A', 'B']]
        tm.assert_frame_equal(res, df.loc[df.A > 2, ['A', 'B']])

        res = df.loc[[2, 3], lambda x: ['A', 'B']]
        tm.assert_frame_equal(res, df.loc[[2, 3], ['A', 'B']])

        res = df.loc[[2, 3], lambda x: ['A', 'B']]
        tm.assert_frame_equal(res, df.loc[[2, 3], ['A', 'B']])

        res = df.loc[3, lambda x: ['A', 'B']]
        tm.assert_series_equal(res, df.loc[3, ['A', 'B']])

        res = df.loc[3, lambda x: ['A', 'B']]
        tm.assert_series_equal(res, df.loc[3, ['A', 'B']])

    def test_frame_loc_callable(self):
        # GH 11485
        df = pd.DataFrame({'X': [1, 2, 3, 4],
                           'Y': list('aabb')},
                          index=list('ABCD'))

        # return label
        res = df.loc[lambda x: ['A', 'C']]
        tm.assert_frame_equal(res, df.loc[['A', 'C']])

        res = df.loc[lambda x: ['A', 'C'], ]
        tm.assert_frame_equal(res, df.loc[['A', 'C'], ])

        res = df.loc[lambda x: ['A', 'C'], :]
        tm.assert_frame_equal(res, df.loc[['A', 'C'], :])

        res = df.loc[lambda x: ['A', 'C'], lambda x: 'X']
        tm.assert_series_equal(res, df.loc[['A', 'C'], 'X'])

        res = df.loc[lambda x: ['A', 'C'], lambda x: ['X']]
        tm.assert_frame_equal(res, df.loc[['A', 'C'], ['X']])

        # mixture
        res = df.loc[['A', 'C'], lambda x: 'X']
        tm.assert_series_equal(res, df.loc[['A', 'C'], 'X'])

        res = df.loc[['A', 'C'], lambda x: ['X']]
        tm.assert_frame_equal(res, df.loc[['A', 'C'], ['X']])

        res = df.loc[lambda x: ['A', 'C'], 'X']
        tm.assert_series_equal(res, df.loc[['A', 'C'], 'X'])

        res = df.loc[lambda x: ['A', 'C'], ['X']]
        tm.assert_frame_equal(res, df.loc[['A', 'C'], ['X']])

    def test_frame_loc_callable_setitem(self):
        # GH 11485
        df = pd.DataFrame({'X': [1, 2, 3, 4],
                           'Y': list('aabb')},
                          index=list('ABCD'))

        # return label
        res = df.copy()
        res.loc[lambda x: ['A', 'C']] = -20
        exp = df.copy()
        exp.loc[['A', 'C']] = -20
        tm.assert_frame_equal(res, exp)

        res = df.copy()
        res.loc[lambda x: ['A', 'C'], :] = 20
        exp = df.copy()
        exp.loc[['A', 'C'], :] = 20
        tm.assert_frame_equal(res, exp)

        res = df.copy()
        res.loc[lambda x: ['A', 'C'], lambda x: 'X'] = -1
        exp = df.copy()
        exp.loc[['A', 'C'], 'X'] = -1
        tm.assert_frame_equal(res, exp)

        res = df.copy()
        res.loc[lambda x: ['A', 'C'], lambda x: ['X']] = [5, 10]
        exp = df.copy()
        exp.loc[['A', 'C'], ['X']] = [5, 10]
        tm.assert_frame_equal(res, exp)

        # mixture
        res = df.copy()
        res.loc[['A', 'C'], lambda x: 'X'] = np.array([-1, -2])
        exp = df.copy()
        exp.loc[['A', 'C'], 'X'] = np.array([-1, -2])
        tm.assert_frame_equal(res, exp)

        res = df.copy()
        res.loc[['A', 'C'], lambda x: ['X']] = 10
        exp = df.copy()
        exp.loc[['A', 'C'], ['X']] = 10
        tm.assert_frame_equal(res, exp)

        res = df.copy()
        res.loc[lambda x: ['A', 'C'], 'X'] = -2
        exp = df.copy()
        exp.loc[['A', 'C'], 'X'] = -2
        tm.assert_frame_equal(res, exp)

        res = df.copy()
        res.loc[lambda x: ['A', 'C'], ['X']] = -4
        exp = df.copy()
        exp.loc[['A', 'C'], ['X']] = -4
        tm.assert_frame_equal(res, exp)

    def test_frame_iloc_callable(self):
        # GH 11485
        df = pd.DataFrame({'X': [1, 2, 3, 4],
                           'Y': list('aabb')},
                          index=list('ABCD'))

        # return location
        res = df.iloc[lambda x: [1, 3]]
        tm.assert_frame_equal(res, df.iloc[[1, 3]])

        res = df.iloc[lambda x: [1, 3], :]
        tm.assert_frame_equal(res, df.iloc[[1, 3], :])

        res = df.iloc[lambda x: [1, 3], lambda x: 0]
        tm.assert_series_equal(res, df.iloc[[1, 3], 0])

        res = df.iloc[lambda x: [1, 3], lambda x: [0]]
        tm.assert_frame_equal(res, df.iloc[[1, 3], [0]])

        # mixture
        res = df.iloc[[1, 3], lambda x: 0]
        tm.assert_series_equal(res, df.iloc[[1, 3], 0])

        res = df.iloc[[1, 3], lambda x: [0]]
        tm.assert_frame_equal(res, df.iloc[[1, 3], [0]])

        res = df.iloc[lambda x: [1, 3], 0]
        tm.assert_series_equal(res, df.iloc[[1, 3], 0])

        res = df.iloc[lambda x: [1, 3], [0]]
        tm.assert_frame_equal(res, df.iloc[[1, 3], [0]])

    def test_frame_iloc_callable_setitem(self):
        # GH 11485
        df = pd.DataFrame({'X': [1, 2, 3, 4],
                           'Y': list('aabb')},
                          index=list('ABCD'))

        # return location
        res = df.copy()
        res.iloc[lambda x: [1, 3]] = 0
        exp = df.copy()
        exp.iloc[[1, 3]] = 0
        tm.assert_frame_equal(res, exp)

        res = df.copy()
        res.iloc[lambda x: [1, 3], :] = -1
        exp = df.copy()
        exp.iloc[[1, 3], :] = -1
        tm.assert_frame_equal(res, exp)

        res = df.copy()
        res.iloc[lambda x: [1, 3], lambda x: 0] = 5
        exp = df.copy()
        exp.iloc[[1, 3], 0] = 5
        tm.assert_frame_equal(res, exp)

        res = df.copy()
        res.iloc[lambda x: [1, 3], lambda x: [0]] = 25
        exp = df.copy()
        exp.iloc[[1, 3], [0]] = 25
        tm.assert_frame_equal(res, exp)

        # mixture
        res = df.copy()
        res.iloc[[1, 3], lambda x: 0] = -3
        exp = df.copy()
        exp.iloc[[1, 3], 0] = -3
        tm.assert_frame_equal(res, exp)

        res = df.copy()
        res.iloc[[1, 3], lambda x: [0]] = -5
        exp = df.copy()
        exp.iloc[[1, 3], [0]] = -5
        tm.assert_frame_equal(res, exp)

        res = df.copy()
        res.iloc[lambda x: [1, 3], 0] = 10
        exp = df.copy()
        exp.iloc[[1, 3], 0] = 10
        tm.assert_frame_equal(res, exp)

        res = df.copy()
        res.iloc[lambda x: [1, 3], [0]] = [-5, -5]
        exp = df.copy()
        exp.iloc[[1, 3], [0]] = [-5, -5]
        tm.assert_frame_equal(res, exp)