diff --git a/pandas/core/indexes/datetimes.py b/pandas/core/indexes/datetimes.py index ee9b948a76ac8..b1463f52333a1 100644 --- a/pandas/core/indexes/datetimes.py +++ b/pandas/core/indexes/datetimes.py @@ -18,7 +18,7 @@ from pandas.core.dtypes.common import _NS_DTYPE, is_float, is_integer, is_scalar from pandas.core.dtypes.dtypes import DatetimeTZDtype -from pandas.core.dtypes.missing import isna +from pandas.core.dtypes.missing import is_valid_nat_for_dtype from pandas.core.accessor import delegate_names from pandas.core.arrays.datetimes import ( @@ -677,8 +677,8 @@ def get_loc(self, key, method=None, tolerance=None): ------- loc : int """ - if is_scalar(key) and isna(key): - key = NaT # FIXME: do this systematically + if is_valid_nat_for_dtype(key, self.dtype): + key = NaT if tolerance is not None: # try converting tolerance now, so errors don't get swallowed by diff --git a/pandas/core/indexes/numeric.py b/pandas/core/indexes/numeric.py index 465f21da1278a..14cb061aa012c 100644 --- a/pandas/core/indexes/numeric.py +++ b/pandas/core/indexes/numeric.py @@ -433,11 +433,7 @@ def get_value(self, series: "Series", key): raise InvalidIndexError loc = self.get_loc(key) - if not is_scalar(loc): - return series.iloc[loc] - - new_values = series._values[loc] - return new_values + return self._get_values_for_loc(series, loc) def equals(self, other) -> bool: """ diff --git a/pandas/tests/indexes/datetimes/test_indexing.py b/pandas/tests/indexes/datetimes/test_indexing.py index f3c255d50aba1..26d99da42cd2b 100644 --- a/pandas/tests/indexes/datetimes/test_indexing.py +++ b/pandas/tests/indexes/datetimes/test_indexing.py @@ -773,3 +773,14 @@ def test_get_loc_nat(self): # GH#20464 index = DatetimeIndex(["1/3/2000", "NaT"]) assert index.get_loc(pd.NaT) == 1 + + assert index.get_loc(None) == 1 + + assert index.get_loc(np.nan) == 1 + + assert index.get_loc(pd.NA) == 1 + + assert index.get_loc(np.datetime64("NaT")) == 1 + + with pytest.raises(KeyError, match="NaT"): + index.get_loc(np.timedelta64("NaT")) diff --git a/pandas/tests/indexes/test_numeric.py b/pandas/tests/indexes/test_numeric.py index 12cc51222e6bb..b83ceb1ce699c 100644 --- a/pandas/tests/indexes/test_numeric.py +++ b/pandas/tests/indexes/test_numeric.py @@ -393,6 +393,46 @@ def test_get_loc_missing_nan(self): # listlike/non-hashable raises TypeError idx.get_loc([np.nan]) + @pytest.mark.parametrize( + "vals", + [ + pd.date_range("2016-01-01", periods=3), + pd.timedelta_range("1 Day", periods=3), + ], + ) + def test_lookups_datetimelike_values(self, vals): + # If we have datetime64 or timedelta64 values, make sure they are + # wrappped correctly + ser = pd.Series(vals, index=range(3, 6)) + ser.index = ser.index.astype("float64") + + expected = vals[1] + + result = ser.index.get_value(ser, 4.0) + assert isinstance(result, type(expected)) and result == expected + result = ser.index.get_value(ser, 4) + assert isinstance(result, type(expected)) and result == expected + + result = ser[4.0] + assert isinstance(result, type(expected)) and result == expected + result = ser[4] + assert isinstance(result, type(expected)) and result == expected + + result = ser.loc[4.0] + assert isinstance(result, type(expected)) and result == expected + result = ser.loc[4] + assert isinstance(result, type(expected)) and result == expected + + result = ser.at[4.0] + assert isinstance(result, type(expected)) and result == expected + # Note: ser.at[4] raises ValueError; TODO: should we make this match loc? + + result = ser.iloc[1] + assert isinstance(result, type(expected)) and result == expected + + result = ser.iat[1] + assert isinstance(result, type(expected)) and result == expected + def test_contains_nans(self): i = Float64Index([1.0, 2.0, np.nan]) assert np.nan in i