CODEWITHSUNDEEP
pythonpython-typingpython-polarsaltair
globglogabgalab
globglogabgalab

Reputation: 451

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!

Upvotes: 4

Views: 196

Answers (1)

ignoring_gravity
ignoring_gravity

Reputation: 10563

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:

Upvotes: 4

Related Questions