Merge branch 'master' into release

Author: tmke8

numpy-stubs/__init__.pyi

@@ -511,7 +511,9 @@ def array(object: Union[ndarray[_DType], _NestedList[ndarray[_DType]]]) -> ndarr
 @overload
 def array(object: Tuple) -> ndarray[bool_]: ...
 @overload
-def arange(stop: int, start: int = ..., step: int = ...) -> ndarray[int64]: ...
+def arange(start: int, stop: int = ..., step: int = ...) -> ndarray[int64]: ...
+@overload
+def arange(start: float, stop: float = ..., step: float = ...) -> ndarray[float64]: ...
 @overload
 def arange(range_: int, dtype: Type[_DType]) -> ndarray[_DType]: ...
 @overload
@@ -561,6 +563,12 @@ def fromiter(iterable: Iterator, dytpe: Type[_DType], count: int = ...) -> ndarr
 def fromstring(
     string: str, dtype: Type[_DType] = ..., count: int = ..., sep: str = ...
 ) -> ndarray: ...
+def frombuffer(
+    buffer: Union[bytes, bytearray, memoryview],
+    dtype: Type[_DType] = ...,
+    count: int = ...,
+    offset: int = ...,
+) -> ndarray: ...
 def histogramdd(
     a: ndarray,
     bins: Optional[Union[ndarray, Series, List, int]],
@@ -730,6 +738,14 @@ def isclose(
 def in1d(
     ar1: ndarray[_DType], ar2: ndarray[_DType], assume_unique: bool = ..., inverse: bool = ...
 ) -> ndarray[bool_]: ...
+def interp(
+    x: _ArrayLike,
+    xp: Sequence[float],
+    fp: Sequence[Union[float, complex]],
+    left: Optional[Union[float, complex]] = ...,
+    right: Optional[Union[float, complex]] = ...,
+    period: Optional[float] = ...,
+) -> ndarray: ...
 def isin(element: Sequence[_DType], test_element: _DType) -> ndarray[_DType]: ...
 @overload
 def isnan(x: float64) -> bool: ...
@@ -896,6 +912,8 @@ nancumsum = cumsum
 nanmean = mean
 nanstd = std
 nansum = sum
+nanmin = min
+nanmax = max
 
 #
 # Saving methods

pandas-stubs/core/frame.pyi

@@ -57,27 +57,30 @@ class DataFrame:
     ): ...
     #
     # magic methods
+    def __add__(self, other: float) -> DataFrame: ...
+    def __and__(self, other: DataFrame) -> DataFrame: ...
     def __eq__(self, other: Union[float, Series, DataFrame]) -> DataFrame: ...  # type: ignore
-    def __ne__(self, other: Union[float, Series, DataFrame]) -> DataFrame: ...  # type: ignore
+    def __floordiv__(self, other: float) -> DataFrame: ...
+    def __ge__(self, other: float) -> DataFrame: ...
+    def __getattr__(self, name: _str) -> Series: ...
     @overload
     def __getitem__(self, idx: _str) -> Series: ...
     @overload
     def __getitem__(
         self, idx: Union[Series, DataFrame, List[_str], Index[_str], _np.ndarray[_np.str_]]
     ) -> DataFrame: ...
-    def __getattr__(self, name: _str) -> Series: ...
+    def __gt__(self, other: float) -> DataFrame: ...
     def __iter__(self) -> Iterator: ...
-    def __len__(self) -> int: ...
-    def __setitem__(self, key: Any, value: Any) -> None: ...
     def __le__(self, other: float) -> DataFrame: ...
+    def __len__(self) -> int: ...
     def __lt__(self, other: float) -> DataFrame: ...
-    def __ge__(self, other: float) -> DataFrame: ...
-    def __gt__(self, other: float) -> DataFrame: ...
     def __mul__(self, other: float) -> DataFrame: ...
-    def __floordiv__(self, other: float) -> DataFrame: ...
-    def __and__(self, other: DataFrame) -> DataFrame: ...
-    def __add__(self, other: float) -> DataFrame: ...
+    def __ne__(self, other: Union[float, Series, DataFrame]) -> DataFrame: ...  # type: ignore
     def __or__(self, other: DataFrame) -> DataFrame: ...
+    def __radd__(self, other: float) -> DataFrame: ...
+    def __rsub__(self, other: float) -> DataFrame: ...
+    def __setitem__(self, key: Any, value: Any) -> None: ...
+    def __sub__(self, other: float) -> DataFrame: ...
     #
     # properties
     @property
@@ -149,20 +152,24 @@ class DataFrame:
     def dropna(
         self,
         inplace: Literal[False] = ...,
-        axis: int = ...,
+        axis: Optional[_AxisType] = ...,
         how: _str = ...,
         subset: _ColSubsetType = ...,
     ) -> DataFrame: ...
     @overload
     def dropna(
-        self, inplace: Literal[True], axis: int = ..., how: _str = ..., subset: _ColSubsetType = ...
+        self,
+        inplace: Literal[True],
+        axis: Optional[_AxisType] = ...,
+        how: _str = ...,
+        subset: _ColSubsetType = ...,
     ) -> None: ...
     @overload
     def fillna(
         self,
         value: Union[float, Dict, Series, DataFrame, _str] = ...,
         method: _str = ...,
-        axis: Union[_str, int] = ...,
+        axis: Optional[_AxisType] = ...,
         inplace: Literal[False] = ...,
         limit: int = ...,
         downcast: Dict = ...,
@@ -173,7 +180,7 @@ class DataFrame:
         inplace: Literal[True],
         value: Union[float, Dict, Series, DataFrame, _str] = ...,
         method: _str = ...,
-        axis: Union[_str, int] = ...,
+        axis: Optional[_AxisType] = ...,
         limit: int = ...,
         downcast: Dict = ...,
     ) -> None: ...
@@ -414,6 +421,15 @@ class DataFrame:
     def to_numpy(self) -> _np.ndarray: ...
     @overload
     def to_numpy(self, dtype: Type[_DTypeNp]) -> _np.ndarray[_DTypeNp]: ...
+    def to_parquet(
+        self,
+        path: Union[Path, _str],
+        engine: Literal["auto", "pyarrow", "fastparquet"] = ...,
+        compression: Union[Literal["snappy", "gzip", "brotli"], None] = ...,
+        index: Optional[bool] = ...,
+        partition_colslist: Optional[List[_str]] = ...,
+        **kwargs: Any,
+    ) -> None: ...
     def unique(self) -> DataFrame: ...
     def update(self, other: Union[DataFrame, Series]) -> None: ...
     def where(self, cond: Union[Series, DataFrame, _np.ndarray]) -> DataFrame: ...

pandas-stubs/core/series.pyi

@@ -121,7 +121,36 @@ class Series(Generic[_DType]):
         self, other: Series, method: Literal["pearson", "kendall", "spearman"] = ...
     ) -> float: ...
     def count(self) -> int: ...
-    def fillna(self, value: Union[_DType, Dict, Series, DataFrame]) -> Series: ...
+    @overload
+    def drop(
+        self, labels: Union[_str, List[_str], Index], axis: _AxisType = ..., inplace: bool = ...
+    ) -> Series[_DType]: ...
+    @overload
+    def drop(self, *, index: Union[List[_str], Index]) -> Series[_DType]: ...
+    @overload
+    def drop(self, *, columns: Union[_str, List[_str], Index]) -> Series[_DType]: ...
+    def drop_duplicates(self, keep: Union[_str, bool] = ...) -> Series[_DType]: ...
+    def duplicated(self, keep: Literal["first", "last", False] = ...) -> Series[_DType]: ...
+    @overload
+    def fillna(
+        self,
+        value: Union[_DType, Dict, Series, DataFrame],
+        method: _str = ...,
+        axis: Optional[_AxisType] = ...,
+        inplace: Literal[True] = ...,
+        limit: int = ...,
+        downcast: Dict = ...,
+    ) -> None: ...
+    @overload
+    def fillna(
+        self,
+        value: Union[_DType, Dict, Series, DataFrame],
+        method: _str = ...,
+        axis: Optional[_AxisType] = ...,
+        inplace: Literal[False] = ...,
+        limit: int = ...,
+        downcast: Dict = ...,
+    ) -> Optional[Series]: ...
     def head(self, n: int = ...) -> Series: ...
     def isna(self) -> Series[bool]: ...
     def isnull(self) -> Series[bool]: ...
@@ -167,6 +196,13 @@ class Series(Generic[_DType]):
     def sort_values(
         self, inplace: Optional[Literal[False]] = ..., ascending: bool = ...
     ) -> Series[_DType]: ...
+    def shift(
+        self,
+        periods: int = ...,
+        freq: Optional[_str] = ...,
+        axis: Optional[_AxisType] = ...,
+        fill_value: Optional[object] = ...,
+    ) -> Series[_DType]: ...
     def std(self) -> float: ...
     def sum(self) -> float: ...
     def to_csv(self, filename: Union[Path, _str], index: bool = ...) -> None: ...
@@ -185,6 +221,7 @@ class Series(Generic[_DType]):
     def tolist(self) -> List[_DType]: ...
     def unique(self) -> _np.ndarray: ...
     def update(self, other: Series) -> None: ...
+    def where(self, cond: Union[Series, DataFrame, _np.ndarray]) -> Series[_DType]: ...
     def value_counts(self, normalize: bool = ...) -> Series[_DType]: ...
     @property
     def at(self) -> _AtIndexerSeries[_DType]: ...

tests/numpy_test.py

@@ -235,3 +235,10 @@ def test_save_load_bytes_io() -> None:
     f.seek(0)
     loaded = np.load(f)
     assert loaded == pytest.approx(a)
+
+
+def test_interp() -> None:
+    assert np.interp([0.5], [0.0, 1.0], [0.0, 100.0]) == pytest.approx([50.0])
+    assert np.interp(
+        [1.5], [0.0, 1.0], [0.0, 100.0], left=-1.0, right=999.0, period=None
+    ) == pytest.approx([999.0])

tests/pandas_test.py

@@ -277,6 +277,13 @@ def test_frame_replace() -> None:
     df.replace(r"1", 1, regex=True, inplace=True)
 
 
+# Uncomment once either pyarrow or fastparquet supports Python 3.9
+# def test_to_parquet(tmp_path: Path) -> None:
+#     filename = str(tmp_path / "data.parq")
+#     df = pd.DataFrame([[1, 2], [3, 4]], columns=["a", "b"])
+#     df.to_parquet(filename, engine="auto", compression=None, index=True, partition_cols=["a"])
+
+
 def test_series_rank() -> None:
     s = pd.Series([0, 1])
     s.rank(method="min")