diff --git a/pandas-stubs/core/frame.pyi b/pandas-stubs/core/frame.pyi index 1f8a2648..fe29a596 100644 --- a/pandas-stubs/core/frame.pyi +++ b/pandas-stubs/core/frame.pyi @@ -698,8 +698,22 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): self, expr: _str, *, + parser: Literal["pandas", "python"] = ..., + engine: Literal["python", "numexpr"] | None = ..., + local_dict: dict[_str, Any] | None = ..., + global_dict: dict[_str, Any] | None = ..., + resolvers: list[Mapping] | None = ..., + level: int = ..., + target: object | None = ..., inplace: Literal[True], - **kwargs: Any, # TODO: make more precise https://github.com/pandas-dev/pandas-stubs/issues/1173 + ) -> None: ... + @overload + def query( + self, + expr: _str, + *, + inplace: Literal[True], + **kwargs: Any, ) -> None: ... @overload def query( @@ -707,7 +721,21 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack): expr: _str, *, inplace: Literal[False] = ..., - **kwargs: Any, # TODO: make more precise https://github.com/pandas-dev/pandas-stubs/issues/1173 + parser: Literal["pandas", "python"] = ..., + engine: Literal["python", "numexpr"] | None = ..., + local_dict: dict[_str, Any] | None = ..., + global_dict: dict[_str, Any] | None = ..., + resolvers: list[Mapping] | None = ..., + level: int = ..., + target: object | None = ..., + ) -> Self: ... + @overload + def query( + self, + expr: _str, + *, + inplace: Literal[False] = ..., + **kwargs: Any, ) -> Self: ... @overload def eval(self, expr: _str, *, inplace: Literal[True], **kwargs: Any) -> None: ... diff --git a/tests/test_frame.py b/tests/test_frame.py index 7d4890ee..807088d1 100644 --- a/tests/test_frame.py +++ b/tests/test_frame.py @@ -514,6 +514,33 @@ def test_types_query() -> None: check(assert_type(df.query("col1 % col2 == 0", inplace=True), None), type(None)) +def test_types_query_kwargs() -> None: + df = pd.DataFrame(data={"col1": [1, 2, 3, 4], "col2": [3, 0, 1, 7]}) + check( + assert_type( + df.query("col1 > col2", parser="pandas", engine="numexpr"), pd.DataFrame + ), + pd.DataFrame, + ) + check( + assert_type( + df.query("col1 > col2", parser="pandas", engine="numexpr", inplace=True), + None, + ), + type(None), + ) + kwargs = {"parser": "pandas", "engine": "numexpr"} + check( + assert_type(df.query("col1 > col2", inplace=False, **kwargs), pd.DataFrame), + pd.DataFrame, + ) + + check( + assert_type(df.query("col1 % col2 == 0", inplace=True, **kwargs), None), + type(None), + ) + + def test_types_eval() -> None: df = pd.DataFrame(data={"col1": [1, 2, 3, 4], "col2": [3, 0, 1, 7]}) check(assert_type(df.eval("E = col1 > col2", inplace=True), None), type(None))