Example #1
0
    def test_get_loc(self):
        # GH 12531
        cidx1 = CategoricalIndex(list('abcde'), categories=list('edabc'))
        idx1 = Index(list('abcde'))
        self.assertEqual(cidx1.get_loc('a'), idx1.get_loc('a'))
        self.assertEqual(cidx1.get_loc('e'), idx1.get_loc('e'))

        for i in [cidx1, idx1]:
            with tm.assertRaises(KeyError):
                i.get_loc('NOT-EXIST')

        # non-unique
        cidx2 = CategoricalIndex(list('aacded'), categories=list('edabc'))
        idx2 = Index(list('aacded'))

        # results in bool array
        res = cidx2.get_loc('d')
        self.assert_numpy_array_equal(res, idx2.get_loc('d'))
        self.assert_numpy_array_equal(
            res, np.array([False, False, False, True, False, True]))
        # unique element results in scalar
        res = cidx2.get_loc('e')
        self.assertEqual(res, idx2.get_loc('e'))
        self.assertEqual(res, 4)

        for i in [cidx2, idx2]:
            with tm.assertRaises(KeyError):
                i.get_loc('NOT-EXIST')

        # non-unique, slicable
        cidx3 = CategoricalIndex(list('aabbb'), categories=list('abc'))
        idx3 = Index(list('aabbb'))

        # results in slice
        res = cidx3.get_loc('a')
        self.assertEqual(res, idx3.get_loc('a'))
        self.assertEqual(res, slice(0, 2, None))

        res = cidx3.get_loc('b')
        self.assertEqual(res, idx3.get_loc('b'))
        self.assertEqual(res, slice(2, 5, None))

        for i in [cidx3, idx3]:
            with tm.assertRaises(KeyError):
                i.get_loc('c')
Example #2
0
    def test_get_loc(self):
        # GH 12531
        cidx1 = CategoricalIndex(list('abcde'), categories=list('edabc'))
        idx1 = Index(list('abcde'))
        self.assertEqual(cidx1.get_loc('a'), idx1.get_loc('a'))
        self.assertEqual(cidx1.get_loc('e'), idx1.get_loc('e'))

        for i in [cidx1, idx1]:
            with tm.assertRaises(KeyError):
                i.get_loc('NOT-EXIST')

        # non-unique
        cidx2 = CategoricalIndex(list('aacded'), categories=list('edabc'))
        idx2 = Index(list('aacded'))

        # results in bool array
        res = cidx2.get_loc('d')
        self.assert_numpy_array_equal(res, idx2.get_loc('d'))
        self.assert_numpy_array_equal(res, np.array([False, False, False,
                                                     True, False, True]))
        # unique element results in scalar
        res = cidx2.get_loc('e')
        self.assertEqual(res, idx2.get_loc('e'))
        self.assertEqual(res, 4)

        for i in [cidx2, idx2]:
            with tm.assertRaises(KeyError):
                i.get_loc('NOT-EXIST')

        # non-unique, slicable
        cidx3 = CategoricalIndex(list('aabbb'), categories=list('abc'))
        idx3 = Index(list('aabbb'))

        # results in slice
        res = cidx3.get_loc('a')
        self.assertEqual(res, idx3.get_loc('a'))
        self.assertEqual(res, slice(0, 2, None))

        res = cidx3.get_loc('b')
        self.assertEqual(res, idx3.get_loc('b'))
        self.assertEqual(res, slice(2, 5, None))

        for i in [cidx3, idx3]:
            with tm.assertRaises(KeyError):
                i.get_loc('c')