allow indexing with .loc using Timestamp

Author: Dr-Irv

pandas-stubs/core/frame.pyi

@@ -113,6 +113,7 @@ from pandas._typing import (
     ReplaceMethod,
     Scalar,
     ScalarT,
+    SequenceNotStr,
     SeriesByT,
     SortKind,
     StataDateFormat,
@@ -194,7 +195,11 @@ class _LocIndexerFrame(_LocIndexer):
     def __getitem__(  # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
         self,
         idx: tuple[
-            int | StrLike | tuple[Scalar, ...] | Callable[[DataFrame], ScalarT],
+            int
+            | StrLike
+            | Timestamp
+            | tuple[Scalar, ...]
+            | Callable[[DataFrame], ScalarT],
             int | StrLike | tuple[Scalar, ...],
         ],
     ) -> Scalar: ...
@@ -207,6 +212,7 @@ class _LocIndexerFrame(_LocIndexer):
                 IndexType
                 | MaskType
                 | _IndexSliceTuple
+                | SequenceNotStr[float | str | Timestamp]
                 | Callable[
                     [DataFrame], ScalarT | list[HashableT] | IndexType | MaskType
                 ],

pandas-stubs/core/series.pyi

@@ -139,6 +139,7 @@ from pandas._typing import (
     ReplaceMethod,
     Scalar,
     ScalarT,
+    SequenceNotStr,
     SeriesByT,
     SortKind,
     StrDtypeArg,
@@ -195,8 +196,7 @@ class _LocIndexerSeries(_LocIndexer, Generic[S1]):
         idx: (
             MaskType
             | Index
-            | Sequence[float]
-            | list[str]
+            | SequenceNotStr[float | str | Timestamp]
             | slice
             | _IndexSliceTuple
             | Sequence[_IndexSliceTuple]

tests/test_frame.py

@@ -3767,3 +3767,17 @@ def test_frame_single_slice() -> None:
     check(assert_type(df.loc[:], pd.DataFrame), pd.DataFrame)
 
     df.loc[:] = 1 + df
+
+
+def test_frame_index_timestamp() -> None:
+    # GH 620
+    dt1 = pd.to_datetime("2023-05-01")
+    dt2 = pd.to_datetime("2023-05-02")
+    s = pd.Series([1, 2], index=[dt1, dt2])
+    df = pd.DataFrame(s)
+    # Next result is Series or DataFrame because the index could be a MultiIndex
+    check(assert_type(df.loc[dt1, :], pd.Series | pd.DataFrame), pd.Series)
+    check(assert_type(df.loc[[dt1], :], pd.DataFrame), pd.DataFrame)
+    df2 = pd.DataFrame({"x": s})
+    check(assert_type(df2.loc[dt1, "x"], Scalar), np.integer)
+    check(assert_type(df2.loc[[dt1], "x"], pd.Series), pd.Series, np.integer)

tests/test_series.py

@@ -3496,3 +3496,12 @@ def test_series_empty_dtype() -> None:
     check(assert_type(pd.Series([]), "pd.Series[Any]"), pd.Series)
     # ensure that an empty string does not get matched to Sequence[Never]
     check(assert_type(pd.Series(""), "pd.Series[str]"), pd.Series)
+
+
+def test_series_index_timestamp() -> None:
+    # GH 620
+    dt1 = pd.to_datetime("2023-05-01")
+    dt2 = pd.to_datetime("2023-05-02")
+    s = pd.Series([1, 2], index=[dt1, dt2])
+    check(assert_type(s[dt1], int), np.integer)
+    check(assert_type(s.loc[[dt1]], "pd.Series[int]"), pd.Series, np.integer)