GH1432 Partial resolution for some TODOs (#1492)

loicdiridollou · MarcoGorelli · cmp0xff · web-flow · commit 085653b47754 · 2025-12-01T10:19:44.000+01:00
* GH1432 Partial resolution * GH1432 PR Feedback * type Index methods: putmask, asof, asof_locs, sort_values, get_indexer_non_unique, get_indexer_for, map, get_slice_bound (#1505) * type Index methods: putmask, asof, asof_locs, sort_values, get_indexer_non_unique, get_indexer_for, map, get_slice_bound * remove redundant `IntervalIndexer.get_indexer_non_unique` * slice_locs can return either `np.integer` or `int` * fix slice_locs test * Apply suggestions from code review Co-authored-by: Yi-Fan Wang <cmp0xff@users.noreply.github.com> * np.integer -> np.intp * test more mask types for `putmask` * Apply suggestions from code review Co-authored-by: Yi-Fan Wang <cmp0xff@users.noreply.github.com> * test asof with str --------- Co-authored-by: Yi-Fan Wang <cmp0xff@users.noreply.github.com> * GH1432 PR Feedback * FIX: #1508 numpy 1darrays in tests (#1512) fix: #1508 numpy arrays in tests * GH1484 Upgrade ty (#1509) * GH1484 Upgrade ty * GH1484 Fix formatting * GH1484 Update PR * GH1432 Partial resolution * GH1432 PR Feedback * Update tests/test_timefuncs.py Co-authored-by: Yi-Fan Wang <cmp0xff@users.noreply.github.com> --------- Co-authored-by: Marco Edward Gorelli <33491632+MarcoGorelli@users.noreply.github.com> Co-authored-by: Yi-Fan Wang <cmp0xff@users.noreply.github.com>
diff --git a/pandas-stubs/_libs/tslibs/timestamps.pyi b/pandas-stubs/_libs/tslibs/timestamps.pyi
@@ -32,6 +32,7 @@ from pandas._libs.tslibs import (
     Tick,
     Timedelta,
 )
+from pandas._libs.tslibs.nattype import NaTType
 from pandas._typing import (
     PeriodFrequency,
     ShapeT,
@@ -236,7 +237,7 @@ class Timestamp(datetime, SupportsIndex):
     def __radd__(
         self, other: np_ndarray[ShapeT, np.timedelta64]
     ) -> np_ndarray[ShapeT, np.datetime64]: ...
-    # TODO: pandas-dev/pandas-stubs#1432 test dt64
+    def __rsub__(self, other: datetime | np.datetime64) -> Timedelta: ...
     @overload  # type: ignore[override]
     def __sub__(self, other: datetime | np.datetime64) -> Timedelta: ...
     @overload
@@ -292,7 +293,15 @@ class Timestamp(datetime, SupportsIndex):
     @property
     def asm8(self) -> np.datetime64: ...
     def tz_convert(self, tz: TimeZones) -> Self: ...
-    # TODO: pandas-dev/pandas-stubs#1432 could return NaT?
+    @overload
+    def tz_localize(  # type: ignore[overload-overlap] # pyright: ignore[reportOverlappingOverload]
+        self,
+        tz: TimeZones,
+        ambiguous: _Ambiguous = "raise",
+        *,
+        nonexistent: Literal["NaT"],
+    ) -> Self | NaTType: ...
+    @overload
     def tz_localize(
         self,
         tz: TimeZones,
diff --git a/pandas-stubs/io/orc.pyi b/pandas-stubs/io/orc.pyi
@@ -1,6 +1,8 @@
 from typing import Any
 
+from fsspec.spec import AbstractFileSystem  # pyright: ignore[reportMissingTypeStubs]
 from pandas import DataFrame
+from pyarrow.fs import FileSystem
 
 from pandas._libs.lib import _NoDefaultDoNotUse
 from pandas._typing import (
@@ -14,8 +16,6 @@ def read_orc(
     path: FilePath | ReadBuffer[bytes],
     columns: list[HashableT] | None = None,
     dtype_backend: DtypeBackend | _NoDefaultDoNotUse = "numpy_nullable",
-    # TODO: pandas-dev/pandas-stubs#1432 type with the correct pyarrow types
-    # filesystem: pyarrow.fs.FileSystem | fsspec.spec.AbstractFileSystem
-    filesystem: Any | None = None,
+    filesystem: FileSystem | AbstractFileSystem | None = None,
     **kwargs: Any,
 ) -> DataFrame: ...
diff --git a/pyproject.toml b/pyproject.toml
@@ -70,6 +70,7 @@ beautifulsoup4 = ">=4.14.2"
 html5lib = ">=1.1"
 python-calamine = ">=0.2.0"
 pyarrow-stubs = { version = ">=20.0.0.20250928", python = "<4" }
+fsspec = "^2025.10.0"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]
diff --git a/tests/scalars/test_scalars.py b/tests/scalars/test_scalars.py
@@ -1386,9 +1386,20 @@ def test_timestamp_misc_methods() -> None:
         pd.Timestamp,
     )
     check(
-        assert_type(ts.tz_localize("US/Pacific", nonexistent="NaT"), pd.Timestamp),
+        assert_type(
+            ts.tz_localize("US/Pacific", nonexistent="NaT"), pd.Timestamp | NaTType
+        ),
         pd.Timestamp,
     )
+    check(
+        assert_type(
+            pd.Timestamp(2025, 3, 9, 2, 30, 0).tz_localize(
+                "US/Eastern", nonexistent="NaT"
+            ),
+            pd.Timestamp | NaTType,
+        ),
+        NaTType,
+    )
     check(
         assert_type(ts.tz_localize("US/Pacific", nonexistent="raise"), pd.Timestamp),
         pd.Timestamp,
diff --git a/tests/test_timefuncs.py b/tests/test_timefuncs.py
@@ -95,13 +95,26 @@ def test_types_init() -> None:
 def test_types_arithmetic() -> None:
     ts = pd.to_datetime("2021-03-01")
     ts2 = pd.to_datetime("2021-01-01")
+    ts_np = np.datetime64("2021-01-01")
+    ts_np_time = np.datetime64("2021-01-01 08:00:05")
     delta = pd.to_timedelta("1 day")
 
     check(assert_type(ts - ts2, pd.Timedelta), pd.Timedelta)
+    check(assert_type(ts - ts_np, pd.Timedelta), pd.Timedelta)
+    check(assert_type(ts - ts_np_time, pd.Timedelta), pd.Timedelta)
     check(assert_type(ts + delta, pd.Timestamp), pd.Timestamp)
     check(assert_type(ts - delta, pd.Timestamp), pd.Timestamp)
     check(assert_type(ts - dt.datetime(2021, 1, 3), pd.Timedelta), pd.Timedelta)
 
+    if TYPE_CHECKING_INVALID_USAGE:
+        # TODO: pandas-dev/pandas-stubs#1511 numpy.datetime64.__sub__ gives datetime.timedelta, which has higher priority
+        assert_type(
+            ts_np - ts, dt.timedelta  # pyright: ignore[reportAssertTypeFailure]
+        )
+        assert_type(
+            ts_np_time - ts, dt.timedelta  # pyright: ignore[reportAssertTypeFailure]
+        )
+
 
 def test_types_comparison() -> None:
     ts = pd.to_datetime("2021-03-01")
