Skip to content

Commit

Permalink
add IntervalIndex.get_loc_exact
Browse files Browse the repository at this point in the history
  • Loading branch information
tp committed Feb 9, 2018
1 parent b835127 commit 7590b87
Show file tree
Hide file tree
Showing 4 changed files with 66 additions and 6 deletions.
1 change: 1 addition & 0 deletions doc/source/whatsnew/v0.23.0.txt
Original file line number Diff line number Diff line change
Expand Up @@ -507,6 +507,7 @@ Other API Changes
- Set operations (union, difference...) on :class:`IntervalIndex` with incompatible index types will now raise a ``TypeError`` rather than a ``ValueError`` (:issue:`19329`)
- :class:`DateOffset` objects render more simply, e.g. "<DateOffset: days=1>" instead of "<DateOffset: kwds={'days': 1}>" (:issue:`19403`)
- :func:`pandas.merge` provides a more informative error message when trying to merge on timezone-aware and timezone-naive columns (:issue:`15800`)
- New :meth:`IntervalIndex.get_loc_exact` has been added to find exact Interval matches only (:issue:`19349`)

.. _whatsnew_0230.deprecations:

Expand Down
52 changes: 47 additions & 5 deletions pandas/core/indexes/interval.py
Original file line number Diff line number Diff line change
Expand Up @@ -903,11 +903,8 @@ def _get_loc_only_exact_matches(self, key):
if not self.is_unique:
raise ValueError("cannot index with a slice Interval"
" and a non-unique index")

# TODO: this expands to a tuple index, see if we can
# do better
return Index(self._multiindex.values).get_loc(key)
raise KeyError
return super(IntervalIndex, self)._engine.get_loc(key)
raise KeyError(key)

def _find_non_overlapping_monotonic_bounds(self, key):
if isinstance(key, IntervalMixin):
Expand Down Expand Up @@ -970,6 +967,10 @@ def get_loc(self, key, method=None):
>>> overlapping_index = pd.IntervalIndex([i2, i3])
>>> overlapping_index.get_loc(1.5)
array([0, 1], dtype=int64)
See Also
--------
get_loc_exact : Exact matches only
"""
self._check_method(method)

Expand Down Expand Up @@ -1003,6 +1004,47 @@ def get_loc(self, key, method=None):
else:
return self._engine.get_loc(key)

def get_loc_exact(self, key, method=None):
"""Get integer location, slice or boolean mask for exact
Interval matches only.
Parameters
----------
key : Interval
The label we want to find locations for. Must have type
:class:`Interval`
method : {None}, optional
* default: matches where the label exactly matches a given
:class:`Interval`.
Returns
-------
loc : int if unique index, slice if monotonic index, else mask
Examples
---------
>>> i1, i2 = pd.Interval(0, 1), pd.Interval(1, 2)
>>> index = pd.IntervalIndex([i1, i2])
>>> index.get_loc_exact(i1)
0
If an exact match is not found, a KeyError is raised
>>> index.get_loc_exact(pd.Interval(0.5, 1.5))
KeyError: Interval(0.5, 1.5, closed='right')
If a label is in several locations, you get all the relevant
locations.
>>> index = pd.IntervalIndex([i1, i2, i1])
>>> index.get_loc_exact(i1)
array([ True, False, True], dtype=bool)
See Also
--------
get_loc
"""
return self.astype(object).get_loc(key, method=method)

def get_value(self, series, key):
if com.is_bool_indexer(key):
loc = key
Expand Down
17 changes: 17 additions & 0 deletions pandas/tests/indexes/interval/test_interval.py
Original file line number Diff line number Diff line change
Expand Up @@ -497,6 +497,23 @@ def test_get_loc_interval(self):
pytest.raises(KeyError, self.index.get_loc,
Interval(-1, 0, 'left'))

# To be removed, replaced by test_interval_new.py (see #16316, #16386)
def test_get_loc_exact(self):
# GH19353
assert self.index.get_loc_exact(Interval(0, 1)) == 0
with pytest.raises(KeyError):
self.index.get_loc_exact(1)
with pytest.raises(KeyError):
self.index.get_loc_exact(Interval(0, 1, 'left'))
with pytest.raises(KeyError):
self.index.get_loc_exact(Interval(0, 0.5))
with pytest.raises(KeyError):
self.index.get_loc_exact(Interval(2, 3))
with pytest.raises(KeyError):
self.index.get_loc_exact(Interval(-1, 0, 'left'))
# See #19353#issuecomment-364295029
self.index.get_loc(Interval(0, 1))

# To be removed, replaced by test_interval_new.py (see #16316, #16386)
def test_get_indexer(self):
actual = self.index.get_indexer([-1, 0, 0.5, 1, 1.5, 2, 3])
Expand Down
2 changes: 1 addition & 1 deletion pandas/tests/indexing/interval/test_interval.py
Original file line number Diff line number Diff line change
Expand Up @@ -136,7 +136,7 @@ def test_with_slices(self):
s[Interval(3, 6):]

expected = s.iloc[3:5]
result = s[[Interval(3, 6)]]
result = s[Interval(3, 6)]
tm.assert_series_equal(expected, result)

# slice of scalar with step != 1
Expand Down

0 comments on commit 7590b87

Please sign in to comment.