Skip to content

Commit a970e3f

Browse files
GH1173 Experiment with fuller typing
1 parent 351810b commit a970e3f

File tree

2 files changed

+34
-0
lines changed

2 files changed

+34
-0
lines changed

pandas-stubs/core/frame.pyi

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -676,6 +676,13 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
676676
self,
677677
expr: _str,
678678
*,
679+
parser: Literal["pandas", "python"] = ...,
680+
engine: Literal["python", "numexpr"] | None = ...,
681+
local_dict: dict[_str, Any] | None = ...,
682+
global_dict: dict[_str, Any] | None = ...,
683+
resolvers: list[Mapping] | None = ...,
684+
level: int = ...,
685+
target: object | None = ...,
679686
inplace: Literal[True],
680687
**kwargs: Any, # TODO: make more precise https://github.com/pandas-dev/pandas-stubs/issues/1173
681688
) -> None: ...
@@ -684,6 +691,13 @@ class DataFrame(NDFrame, OpsMixin, _GetItemHack):
684691
self,
685692
expr: _str,
686693
*,
694+
parser: Literal["pandas", "python"] = "pandas",
695+
engine: Literal["python", "numexpr"] | None = "numexpr",
696+
local_dict: dict[_str, Any] | None = None,
697+
global_dict: dict[_str, Any] | None = None,
698+
resolvers: list[Mapping] | None = None,
699+
level: int = 0,
700+
target: object | None = None,
687701
inplace: Literal[False] = ...,
688702
**kwargs: Any, # TODO: make more precise https://github.com/pandas-dev/pandas-stubs/issues/1173
689703
) -> Self: ...

tests/test_frame.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -513,6 +513,26 @@ def test_types_query() -> None:
513513
check(assert_type(df.query("col1 % col2 == 0", inplace=True), None), type(None))
514514

515515

516+
def test_types_query_kwargs() -> None:
517+
df = pd.DataFrame(data={"col1": [1, 2, 3, 4], "col2": [3, 0, 1, 7]})
518+
check(
519+
assert_type(
520+
df.query("col1 > col2", parser="pandas", engine="numexpr"), pd.DataFrame
521+
),
522+
pd.DataFrame,
523+
)
524+
check(
525+
assert_type(
526+
df.query("col1 > col2", parser="pandas", engine="numexpr"), pd.DataFrame
527+
),
528+
pd.DataFrame,
529+
)
530+
kwargs = {"parser": "pandas", "engine": "numexpr"}
531+
check(assert_type(df.query("col1 > col2", **kwargs), pd.DataFrame), pd.DataFrame)
532+
533+
check(assert_type(df.query("col1 % col2 == 0", inplace=True), None), type(None))
534+
535+
516536
def test_types_eval() -> None:
517537
df = pd.DataFrame(data={"col1": [1, 2, 3, 4], "col2": [3, 0, 1, 7]})
518538
check(assert_type(df.eval("E = col1 > col2", inplace=True), None), type(None))

0 commit comments

Comments
 (0)