GH1173 Experiment with fuller typing

Author: loicdiridollou

pandas-stubs/core/frame.pyi

@@ -676,6 +676,13 @@ 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: ...
@@ -684,6 +691,13 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
         self,
         expr: _str,
         *,
+        parser: Literal["pandas", "python"] = "pandas",
+        engine: Literal["python", "numexpr"] | None = "numexpr",
+        local_dict: dict[_str, Any] | None = None,
+        global_dict: dict[_str, Any] | None = None,
+        resolvers: list[Mapping] | None = None,
+        level: int = 0,
+        target: object | None = None,
         inplace: Literal[False] = ...,
         **kwargs: Any,  # TODO: make more precise https://github.com/pandas-dev/pandas-stubs/issues/1173
     ) -> Self: ...

tests/test_frame.py

@@ -513,6 +513,26 @@ 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"), pd.DataFrame
+        ),
+        pd.DataFrame,
+    )
+    kwargs = {"parser": "pandas", "engine": "numexpr"}
+    check(assert_type(df.query("col1 > col2", **kwargs), pd.DataFrame), pd.DataFrame)
+
+    check(assert_type(df.query("col1 % col2 == 0", inplace=True), 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))