search

pola-rs / polars

Dataframes powered by a multithreaded, vectorized query engine, written in Rust
https://docs.pola.rs
Other
29.58k stars 1.89k forks source link

Type hint support for `GroupBy` #15459

Open dave-edison opened 6 months ago

dave-edison commented 6 months ago

Description

I have a function which returns a GroupBy, but it doesn't seem possible to provide type hints for it in a way that appeases static type checkers (or at least pyright, can't speak to mypy or others).

Using polars.dataframe.group_by.GroupBy gives an reportAttributeAccessIssue error, and I don't see anything in polars.type_aliases that would work here.

Certainly possible I'm missing something! If not, it would be nice to be able to give hints for this, although I also understand that returning a GroupBy in a function is somewhat non-standard usage.

runfalk commented 3 months ago

Not sure if this is still a problem, with pyright 1.1.367 this code seem to work for me:

import polars as pl
from polars.dataframe.group_by import GroupBy
def f() -> GroupBy:
    return pl.DataFrame().group_by("a")
t.reveal_type(f())  # GroupBy as expected

This is in strict mode.

[tool.pyright]
typeCheckingMode = "strict"
useLibraryCodeForTypes = true
pythonVersion = "3.12"
reportUnnecessaryTypeIgnoreComment = true
reportMissingTypeStubs = false