如何为Pandas与Pandas-on-Spark DataFrame的联合类型添加类型提示？

问题核心是pandas和pyspark.pandas的方法返回类型注解不统一，导致mypy无法正确推断泛型类型。以下是几种实用的解决方法：

方案1：自定义协议（Protocol）统一接口 #

通过定义协议，明确我们需要的DataFrame方法签名，强制loc和rename返回与自身一致的类型，让mypy能正确识别泛型。

import pyspark.pandas as ps
import pandas as pd
from typing import TypeVar, Protocol, Any

class DataFrameLike(Protocol):
    def loc(self, row: Any, cols: list[str]) -> "DataFrameLike": ...
    def rename(self, columns: dict[str, str]) -> "DataFrameLike": ...

ADataFrame = TypeVar("ADataFrame", bound=DataFrameLike)

def format_table(df: ADataFrame) -> ADataFrame:
    renamed = df.loc[:, ["ID", "code"]].rename(columns={"ID": "new_id", "code": "new_code"})
    return renamed

方案2：类型守卫+分支处理 #

直接区分输入类型，用类型守卫让mypy明确每个分支的输入和返回类型，消除类型歧义。

import pyspark.pandas as ps
import pandas as pd
from typing import Union, TypeGuard

def is_pandas_df(df: Union[pd.DataFrame, ps.DataFrame]) -> TypeGuard[pd.DataFrame]:
    return isinstance(df, pd.DataFrame)

def is_ps_df(df: Union[pd.DataFrame, ps.DataFrame]) -> TypeGuard[ps.DataFrame]:
    return isinstance(df, ps.DataFrame)

def format_table(df: Union[pd.DataFrame, ps.DataFrame]) -> Union[pd.DataFrame, ps.DataFrame]:
    if is_pandas_df(df):
        return df.loc[:, ["ID", "code"]].rename(columns={"ID": "new_id", "code": "new_code"})
    elif is_ps_df(df):
        return df.loc[:, ["ID", "code"]].rename(columns={"ID": "new_id", "code": "new_code"})
    raise TypeError("Unsupported DataFrame type")

方案3：统一转换类型处理（适合小数据场景） #

把输入统一转成pyspark.pandas DataFrame处理，最后再转回原类型，代价是可能有数据转换的性能开销。

import pyspark.pandas as ps
import pandas as pd
from typing import Union

def format_table(df: Union[pd.DataFrame, ps.DataFrame]) -> Union[pd.DataFrame, ps.DataFrame]:
    is_pandas_input = isinstance(df, pd.DataFrame)
    # 统一转成ps.DataFrame处理
    processed_df = ps.from_pandas(df) if is_pandas_input else df
    renamed_df = processed_df.loc[:, ["ID", "code"]].rename(columns={"ID": "new_id", "code": "new_code"})
    # 转回原类型返回
    return renamed_df.to_pandas() if is_pandas_input else renamed_df

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