I use polars
dataframes (new to this module) and I'm using some static typing, to keep my code tidy and clean for debugging purposes, and to allow auto-completion of methods and attributes on my editor. Everything goes well.
However, when plotting things from dataframes with the altair
API, as shown in the doc, I am unable to find the type of the returned object in polars
.
import polars as pl
import typing as tp
data = {"a": [0,0,0,0,1,1,1,2,2,3]}
df = pl.DataFrame(data)
def my_plot(df: pl.DataFrame, col: str) -> tp.Any:
"""Plot an histogram of the distribution of values of df[col]"""
return df[col].value_counts(
).plot.bar(
y="count",
x=col
).properties(
width=400,
)
u = my_plot(df, "a")
u.show()
How do I type the output of this function? The doc states the output of (...).plot
is a DataFramePlot
object but there is no info on this type, and anyway I'm using the output of (...).plot.bar(...)
which has a different type.
If I run type(u)
, I get altair.vegalite.v5.api.Chart
but it seems sketchy to use this for static typing, and I don't want to import altair
in my code as the altair
methods I need are already included in polars
.
I couldn't find any info about this so any help is welcome
Thanks!
You can do
from __future__ import annotations
from typing import TYPE_CHECKING
if TYPE_CHECKING:
import altair as alt
def my_plot(df: pl.DataFrame, col: str) -> alt.Chart: