Skip to content

ENH: Type support for variables in DataFrame.query() #61731

Closed
@malekzada

Description

@malekzada

Feature Type

  • Adding new functionality to pandas

  • Changing existing functionality in pandas

  • Removing existing functionality in pandas

Problem Description

Now using variables inside df.query("col > @my_var") doesn’t produce strong typing:
IDEs don’t catch type mismatches (e.g. my_var is a string but col is numeric).

Feature Description

Add type support in pandas-stubs so that functions like query():

  • Accept variables bound via @
  • Validate that their types align with the DataFrame column dtype
  • Offer autocomplete in IDEs

Example:

from typing import TypedDict

class Record(TypedDict):
    a: int
    b: str

df: DataFrame[Record] = ...
my_var: int = 5
filtered = df.query("a > @my_var") 

other_var: str = "foo"
df.query("a > @other_var")  
# Should flag type mismatch in IDE/type-checker

### Alternative Solutions

```python
# Validate variable type before calling query
from typing import assert_type

my_var = 5
assert_type(my_var, int)  # Mypy will enforce this
df.query("a > @my_var")

OR

# Type-safe alternative using boolean indexing
df[df["a"] > my_var]  # Fully type-checkable, no strings

### Additional Context

_No response_

Metadata

Metadata

Assignees

No one assigned

    Labels

    EnhancementNeeds TriageIssue that has not been reviewed by a pandas team member

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions