gh-591: changed the return type of argmin and argmax (#609)

* changed the return type and added test

* changed the return type

Author: ramvikrams

pandas-stubs/core/base.pyi

@@ -57,10 +57,10 @@ class IndexOpsMixin(OpsMixin):
     def min(self, axis=..., skipna: bool = ..., **kwargs): ...
     def argmax(
         self, axis: AxisIndex | None = ..., skipna: bool = ..., *args, **kwargs
-    ) -> np.ndarray: ...
+    ) -> np.int64: ...
     def argmin(
         self, axis: AxisIndex | None = ..., skipna: bool = ..., *args, **kwargs
-    ) -> np.ndarray: ...
+    ) -> np.int64: ...
     def tolist(self) -> list: ...
     def to_list(self) -> list: ...
     def __iter__(self): ...

tests/test_pandas.py

@@ -1973,3 +1973,11 @@ def g(x: pd.Series) -> int:
         ),
         pd.DataFrame,
     )
+
+
+def test_argmin_and_argmax_return() -> None:
+    df = pd.DataFrame({"a": [-1, 0, 1], "b": [1, 2, 3]})
+    i = df.a.abs().argmin()
+    i1 = df.a.abs().argmax()
+    check(assert_type(i, np.int64), np.int64)
+    check(assert_type(i1, np.int64), np.int64)