如何避免链式when/then赋值中的Mypy不兼容类型警告？

问题根源在于初始代码中pl.when(False).then(None)直接传入了Python原生的None，而非Polars的Expr类型，导致Mypy无法正确推断后续链式调用的类型，引发不兼容警告。以下是两种解决方法：

方法一：修正初始值为Polars Expr类型 #

将初始的then(None)替换为then(pl.lit(None))，确保整个链式调用的类型统一为Polars的WhenThenBuilder：

import polars as pl

expr = pl.when(False).then(pl.lit(None))

for pattern, replacement in replacement_rules.items():
    expr = expr.when(pl.col("data").str.contains(pattern)).then(pl.lit(replacement))
   
expr = expr.when(pl.col("ISO_codes").str.len_chars() > 0).then(            
    pl.col("ISO_codes")
      .replace(iso_translation, default="Unknown ISO Code")
)

方法二：添加明确的类型注解 #

如果修正初始值后仍有警告，可以给expr添加精确的类型注解，帮助Mypy正确推断类型：

import polars as pl
from typing import cast

# 显式指定为WhenThenBuilder类型
expr: pl.WhenThenBuilder = cast(pl.WhenThenBuilder, pl.when(False).then(pl.lit(None)))

for pattern, replacement in replacement_rules.items():
    expr = expr.when(pl.col("data").str.contains(pattern)).then(pl.lit(replacement))
   
expr = expr.when(pl.col("ISO_codes").str.len_chars() > 0).then(            
    pl.col("ISO_codes")
      .replace(iso_translation, default="Unknown ISO Code")
)

额外优化：合并链式调用 #

可以将循环内的when和then调用合并为一行，让代码更简洁，同时也有助于Mypy的类型推断：

import polars as pl

expr = pl.when(False).then(pl.lit(None))

for pattern, replacement in replacement_rules.items():
    expr = expr.when(pl.col("data").str.contains(pattern)).then(pl.lit(replacement))
   
expr = expr.when(pl.col("ISO_codes").str.len_chars() > 0).then(
    pl.col("ISO_codes").replace(iso_translation, default="Unknown ISO Code")
)

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