Mypy报union-attr错误：Optional[SparkSession]的None项无createDataFrame属性

环境版本 #

• pyspark==3.2.2
• mypy == 0.961
• python == 3.8.1

测试代码 #

class TestParsingFunctions(unittest.TestCase):
    """
    该类用于测试mintel_import.py中的函数
    """

    db_name = None
    spark = None

    @classmethod
    def setUpClass(cls):
        cls.test_dir = tempfile.TemporaryDirectory().name
        findspark.init()
        cls.spark = SparkSession.builder.master("local[1]").getOrCreate()

    def test_df_remove_active_ingredient_pattern(self):
        input_df = self.spark.createDataFrame(
            data=[
                (
                    "9253090",
                    "This input text is TRANSFORMED",
                ),
                
            ],
            schema="ID: string, INPUT_TEXT: string",
        )

        expected_result = self.spark.createDataFrame(
            data=[
                ("9253090", "thisinputtextistransformed"),
                
            ],
            schema="ID: string, INPUT_TEXT: string",
        )
        result_df = transfor_text(
            input_df, col_name="INPUT_TEXT"
        )
        rows = result_df.sort("ID").collect()
        expected_rows = expected_result.sort("ID").collect()
        # 逐行比较DataFrame
        for row_num, row in enumerate(rows):
            assert row == expected_rows[row_num], f"Error on row: {row_num}"

if __name__ == "__main__":
    unittest.main()

Mypy报错信息 #

错误："Optional[SparkSession]"中的"None"项没有属性"createDataFrame" [union-attr]

解决方法 #

报错原因是类属性spark初始化为None，Mypy静态检查认为它可能始终为None，无法确认调用createDataFrame的安全性。以下是几种修复方式：

方法1：明确类型注解 #

给spark属性添加SparkSession类型注解，直接声明它最终会是有效实例：

from pyspark.sql import SparkSession

class TestParsingFunctions(unittest.TestCase):
    db_name = None
    spark: SparkSession  # 明确类型，排除Optional可能性

    @classmethod
    def setUpClass(cls):
        cls.test_dir = tempfile.TemporaryDirectory().name
        findspark.init()
        cls.spark = SparkSession.builder.master("local[1]").getOrCreate()
    
    # 后续测试方法保持不变

方法2：添加非空断言 #

在测试方法开头添加断言，告诉Mypyself.spark不可能为None：

def test_df_remove_active_ingredient_pattern(self):
    assert self.spark is not None, "SparkSession未完成初始化"
    input_df = self.spark.createDataFrame(
        # 原有代码内容
    )

方法3：临时忽略类型检查 #

如果需要快速绕过检查，可以在调用createDataFrame的行添加忽略注释：

input_df = self.spark.createDataFrame(  # type: ignore[union-attr]
    # 原有代码内容
)

优先推荐方法1，从类型定义层面彻底解决问题，代码可读性更强；方法2适合临时验证场景；方法3属于应急方案，不建议长期使用。

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