Fix for inplace methods, tests (#22)

Author: zkrolikowski-vl

setup.py

@@ -21,7 +21,7 @@ def list_packages(source_path: str = src_path) -> None:
 setup(
     name="pandas-stubs",
     package_dir={"": src_path},
-    version="1.0.4.4",
+    version="1.0.4.5",
     description="Type annotations for Pandas",
     long_description=(open("README.md").read()
                       if os.path.exists("README.md") else ""),

tests/snippets/test_frame.py

@@ -111,6 +111,22 @@ def test_types_sort_index() -> None:
     df.sort_index(kind="mergesort")
 
 
+def test_types_set_index() -> None:
+    df = pd.DataFrame(data={'col1': [1, 2, 3, 4], 'col2': ['a', 'b', 'c', 'd']}, index=[5, 1, 3, 2])
+    df.set_index('col1')
+    df.set_index('col1', drop=False)
+    df.set_index('col1', append=True)
+    df.set_index('col1', verify_integrity=True)
+    df.set_index(['col1', 'col2'])
+    df.set_index('col1', inplace=True)
+
+
+def test_types_query() -> None:
+    df = pd.DataFrame(data={'col1': [1, 2, 3, 4], 'col2': [3, 0, 1, 7]})
+    df.query("col1 > col2")
+    df.query("col1 % col2 == 0", inplace=True)
+
+
 def test_types_sort_values() -> None:
     df = pd.DataFrame(data={'col1': [2, 1], 'col2': [3, 4]})
     df.sort_values('col1')

third_party/3/pandas/core/frame.pyi

@@ -5,7 +5,7 @@ import numpy.ma as np
 from pandas import datetime
 from pandas._typing import Axes, Axis, Dtype as Dtype, FilePathOrBuffer, Level, Renamer, Column, Label, FrameOrSeries, \
     ArrayLike, AnyArrayLike, GoogleCredentials, Scalar, ReplaceMethod, ToReplace, ReplaceValue, Frequency, AxisOption, \
-    Orientation, Function, AggregationFunction, GroupByObject, GeneralDuplicatesKeepStrategy, InterpolationMethod, NamedCorrelationMethod, \
+    Orientation, Function, AggregationFunction, GroupByObject, GeneralDuplicatesKeepStrategy, InterpolationMethod, \
     CorrelationMethod, SortKind, JoinType, FillMethod, ErrorsStrategy, NaSortPosition, FillValue, TimestampMethod
 from pandas.core.accessor import CachedAccessor
 from pandas.core.base import PandasObject
@@ -118,15 +118,15 @@ class DataFrame(NDFrame):
     def transpose(self, *args: Any, copy: bool = ...) -> DataFrame: ...
     @property
     def T(self) -> DataFrame: ...
-    def query(self, expr: str, inplace: bool = ..., **kwargs: Any) -> DataFrame: ...
-    def eval(self, expr: str, inplace: bool = ..., **kwargs: Any) -> Union[np.ndarray, int, float, PandasObject]: ...
+    def query(self, expr: str, inplace: bool = ..., **kwargs: Any) -> Optional[DataFrame]: ...
+    def eval(self, expr: str, inplace: bool = ..., **kwargs: Any) -> Union[None, np.ndarray, int, float, PandasObject]: ...
     def select_dtypes(self, include: Optional[Sequence[Union[str, Dtype]]] = ..., exclude: Optional[Sequence[Union[str, Dtype]]] = ...) -> DataFrame: ...
     def insert(self, loc: int, column: Union[Column, Hashable], value: Union[int, Series, ArrayLike], allow_duplicates: Optional[bool] = ...) -> None: ...
     def assign(self, **kwargs: Any) -> DataFrame: ...
     def lookup(self, row_labels: Sequence[Any], col_labels: Sequence[Column]) -> np.ndarray: ...
     def align(self, other: FrameOrSeries, join: JoinType = ..., axis: AxisOption = ..., level: Level = ..., copy: bool = ..., fill_value: Scalar = ..., method: Optional[FillMethod] = ..., limit: Optional[int] = ..., fill_axis: AxisOption = ..., broadcast_axis: AxisOption = ...) -> DataFrame: ...
     def reindex(self, *args: Any, **kwargs: Any) -> DataFrame: ...
-    def drop(self, labels: Optional[Sequence[Label]] = ..., axis: AxisOption = ..., index: Optional[Sequence[Label]] = ..., columns: Optional[Sequence[Label]] = ..., level: Optional[Level] = ..., inplace: bool = ..., errors: ErrorsStrategy = ...) -> DataFrame: ...
+    def drop(self, labels: Optional[Sequence[Label]] = ..., axis: AxisOption = ..., index: Optional[Sequence[Label]] = ..., columns: Optional[Sequence[Label]] = ..., level: Optional[Level] = ..., inplace: bool = ..., errors: ErrorsStrategy = ...) -> Optional[DataFrame]: ...
     def rename(self, mapper: Optional[Renamer] = ..., *, index: Optional[Renamer] = ..., columns: Optional[Renamer] = ..., axis: Optional[Axis] = ..., copy: bool = ..., inplace: bool = ..., level: Optional[Level] = ..., errors: ErrorsStrategy = ...) -> Optional[DataFrame]: ...
     def fillna(self, value: FillValue = ..., method: Optional[FillMethod] = ..., axis: Optional[Axis] = ..., inplace: Optional[bool] = ..., limit: int = ..., downcast: Optional[Dict[Any, Dtype]] = ...) -> Optional[DataFrame]: ...
     def replace(self, to_replace: Optional[ToReplace] = ..., value: Optional[ReplaceValue] = ..., inplace: bool = ..., limit: Optional[int] = ..., regex: bool = ..., method: ReplaceMethod = ...) -> DataFrame: ...
@@ -137,7 +137,7 @@ class DataFrame(NDFrame):
     def isnull(self) -> DataFrame: ...
     def notna(self) -> DataFrame: ...
     def notnull(self) -> DataFrame: ...
-    def dropna(self, axis: AxisOption = ..., how: DropTypes = ..., thresh: Optional[int] = ..., subset: Optional[Any] = ..., inplace: bool = ...) -> DataFrame: ...
+    def dropna(self, axis: AxisOption = ..., how: DropTypes = ..., thresh: Optional[int] = ..., subset: Optional[Any] = ..., inplace: bool = ...) -> Optional[DataFrame]: ...
     def drop_duplicates(self, subset: Optional[Union[Hashable, Sequence[Hashable]]] = ..., keep: GeneralDuplicatesKeepStrategy = ..., inplace: bool = ..., ignore_index: bool = ...) -> Optional[DataFrame]: ...
     def duplicated(self, subset: Optional[Union[Hashable, Sequence[Hashable]]] = ..., keep: Union[str, bool] = ...) -> Series: ...
     # Parent allowed by to be None - that's the reason for override

third_party/3/pandas/core/generic.pyi

@@ -58,7 +58,7 @@ class NDFrame(PandasObject, SelectionMixin, indexing.IndexingMixin):
     def squeeze(self, axis: Optional[Any] = ...) -> Any: ...
     def swaplevel(self, i: Any = ..., j: Any = ..., axis: Any = ...) -> FrameOrSeries: ...
     def rename(self, mapper: Optional[Renamer]=..., *, index: Optional[Renamer]=..., columns: Optional[Renamer]=..., axis: Optional[Axis]=..., copy: bool_t=..., inplace: bool_t=..., level: Optional[Level]=..., errors: ErrorsStrategy=...) -> Optional[FrameOrSeries]: ...
-    def rename_axis(self, mapper: Any = ..., **kwargs: Any) -> Any: ...
+    def rename_axis(self, mapper: Any = ..., **kwargs: Any) -> Optional[FrameOrSeries]: ...
     def equals(self, other: Any) -> Any: ...
     @property
     def empty(self) -> bool_t: ...
@@ -98,14 +98,14 @@ class NDFrame(PandasObject, SelectionMixin, indexing.IndexingMixin):
     def fillna(self, value: Any = ..., method: Any = ..., axis: Any = ..., inplace: bool_t=..., limit: Any = ..., downcast: Any = ...) -> Optional[FrameOrSeries]: ...
     def ffill(self, axis: Any = ..., inplace: bool_t=..., limit: Any = ..., downcast: Any = ...) -> Optional[FrameOrSeries]: ...
     def bfill(self, axis: Any = ..., inplace: bool_t=..., limit: Any = ..., downcast: Any = ...) -> Optional[FrameOrSeries]: ...
-    def replace(self, to_replace: Optional[ToReplace] = ..., value: Optional[ReplaceValue] = ..., inplace: bool_t = ..., limit: Optional[int] = ..., regex: bool_t = ..., method: ReplaceMethod = ...) -> FrameOrSeries: ...
-    def interpolate(self, method: str = ..., axis: int = ..., limit: Optional[Any] = ..., inplace: bool_t = ..., limit_direction: str = ..., limit_area: Optional[Any] = ..., downcast: Optional[Any] = ..., **kwargs: Any) -> Any: ...
+    def replace(self, to_replace: Optional[ToReplace] = ..., value: Optional[ReplaceValue] = ..., inplace: bool_t = ..., limit: Optional[int] = ..., regex: bool_t = ..., method: ReplaceMethod = ...) -> Optional[FrameOrSeries]: ...
+    def interpolate(self, method: str = ..., axis: int = ..., limit: Optional[Any] = ..., inplace: bool_t = ..., limit_direction: str = ..., limit_area: Optional[Any] = ..., downcast: Optional[Any] = ..., **kwargs: Any) -> Optional[FrameOrSeries]: ...
     def asof(self, where: Any, subset: Optional[Any] = ...) -> Any: ...
     def isna(self) -> FrameOrSeries: ...
     def isnull(self) -> FrameOrSeries: ...
     def notna(self) -> FrameOrSeries: ...
     def notnull(self) -> FrameOrSeries: ...
-    def clip(self, lower: Any = ..., upper: Any = ..., axis: Any = ..., inplace: bool_t=..., *args: Any, **kwargs: Any) -> FrameOrSeries: ...
+    def clip(self, lower: Any = ..., upper: Any = ..., axis: Any = ..., inplace: bool_t=..., *args: Any, **kwargs: Any) -> Optional[FrameOrSeries]: ...
     def asfreq(self, freq: Any, method: Any = ..., how: Optional[str]=..., normalize: bool_t=..., fill_value: Any = ...) -> FrameOrSeries: ...
     def at_time(self, time: Any, asof: bool_t=..., axis: Any = ...) -> FrameOrSeries: ...
     def between_time(self, start_time: Any, end_time: Any, include_start: bool_t=..., include_end: bool_t=..., axis: Any = ...) -> FrameOrSeries: ...