如何为Python中Polars的Altair绘图返回值添加类型注解？

方案1：字符串类型标注（Python 3.10+） #

直接用字符串形式指定返回类型，无需导入Altair，类型检查器和编辑器可正确识别该类型并提供补全支持：

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) -> "altair.vegalite.v5.api.Chart":
    """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()

方案2：条件导入（兼容所有Python版本） #

利用typing.TYPE_CHECKING常量，仅在类型检查阶段导入Altair，运行时不会执行导入操作，既满足静态检查需求，又避免运行时引入额外依赖：

import polars as pl
import typing as tp
from typing import TYPE_CHECKING

# 仅类型检查时导入，运行时不生效
if TYPE_CHECKING:
    import altair.vegalite.v5.api as altair_api

data = {"a": [0,0,0,0,1,1,1,2,2,3]}

df = pl.DataFrame(data)

def my_plot(df: pl.DataFrame, col: str) -> altair_api.Chart:
    """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()

方案3：临时过渡方案——使用tp.Unknown（Python 3.10+） #

如果暂时不想明确指定具体类型，又想比tp.Any更严谨，可以用tp.Unknown，它表示类型未知但非任意类型，类型检查器会对后续方法调用做基本的存在性校验：

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.Unknown:
    """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()

补充说明 #

Polars的plot方法返回的DataFramePlot是一个封装类，链式调用bar等绘图方法后，底层会直接返回Altair的Chart对象，这也是type(u)返回altair.vegalite.v5.api.Chart的原因。Polars本身未为该返回类型提供自定义包装，因此必须直接引用Altair的类型来实现精确标注。

内容的提问来源于stack exchange，提问作者globglogabgalab