PySpark中Mypy为何混淆Sequence[str]与Union[str,List[str]]？

问题重现 #

你编写的函数根据布尔参数切换分区列的数量，但Mypy抛出类型不兼容错误：

from pyspark.sql import Window, WindowSpec

def get_window(single_column: bool) -> WindowSpec:
    partition_cols = "key" if single_column else ["key", "name"]
    return Window.partitionBy(partition_cols).orderBy("timestamp").rangeBetween(0, 10)

Mypy错误输出：

$ mypy tmp.py
tmp.py:8: error: Argument 1 to "partitionBy" of "Window" has incompatible type "Sequence[str]"; expected "Union[Union[Column, str], List[Union[Column, str]]]"  [arg-type]

问题原因 #

Mypy会将三元表达式"key" if single_column else ["key", "name"]的类型推断为Sequence[str]，而PySpark的Window.partitionBy参数类型定义为Union[Union[Column, str], List[Union[Column, str]]]——这里要求的是具体的List类型，而非更宽泛的Sequence，因此触发类型不匹配错误。

解决方案 #

方案1：显式标注变量类型

给partition_cols指定明确的Union[str, List[str]]类型，让Mypy正确识别变量的可能类型：

from pyspark.sql import Window, WindowSpec
from typing import Union, List

def get_window(single_column: bool) -> WindowSpec:
    partition_cols: Union[str, List[str]] = "key" if single_column else ["key", "name"]
    return Window.partitionBy(partition_cols).orderBy("timestamp").rangeBetween(0, 10)

方案2：拆分分支处理

直接拆分if-else分支，分别调用partitionBy，避免类型推断歧义：

from pyspark.sql import Window, WindowSpec

def get_window(single_column: bool) -> WindowSpec:
    if single_column:
        return Window.partitionBy("key").orderBy("timestamp").rangeBetween(0, 10)
    else:
        return Window.partitionBy(["key", "name"]).orderBy("timestamp").rangeBetween(0, 10)

方案3：临时忽略类型错误（不推荐）

如果只是临时绕过检查，可以添加类型忽略注释，但不建议长期使用，会掩盖潜在问题：

from pyspark.sql import Window, WindowSpec

def get_window(single_column: bool) -> WindowSpec:
    partition_cols = "key" if single_column else ["key", "name"]
    return Window.partitionBy(partition_cols).orderBy("timestamp").rangeBetween(0, 10)  # type: ignore[arg-type]

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