pythonpython-typingpython-polarsaltair

How to type Polars' Altair plots in Python?


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!


Solution

  • 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: