如何在PySpark中较好地运用PEP 484类型标注？

以下是几种可行的解决方案，针对PySpark与JVM交互时PEP 484类型标注失效的问题：

1. 自定义类型存根文件 #

为JVM导出的Hadoop类创建.pyi存根文件，手动定义类型接口，让类型检查器（如mypy）识别这些JVM对象的结构：

# hadoop_stubs.pyi
from typing import Any

class Path:
    def __init__(self, path_str: str) -> None: ...
    def getFileSystem(self, conf: Any) -> FileSystem: ...

class FileSystem:
    def exists(self, path: Path) -> bool: ...
    def delete(self, path: Path, recursive: bool) -> bool: ...

在业务代码中导入存根类型并标注：

from typing import Callable
from hadoop_stubs import Path, FileSystem

class SparkAsset:
    def __init__(self, client):
        self._client = client
        self.path = "/some/hdfs/path"

    @property
    def hadoop_configuration(self) -> Any:
        return self._client.sparkContext._jsc.hadoopConfiguration()

    @property
    def filesystem_path_func(self) -> Callable[[str], Path]:
        return self._client.sparkContext._jvm.org.apache.hadoop.fs.Path

    @property
    def filesystem_path(self) -> Path:
        return self.filesystem_path_func(self.path)

    @property
    def filesystem(self) -> FileSystem:
        return self.filesystem_path.getFileSystem(self.hadoop_configuration)

存根文件无需实现逻辑，仅需定义类型结构即可辅助类型检查。

2. 使用类型断言替代# type: ignore #

通过typing.cast明确指定JVM对象的类型，让类型检查器认可标注：

from typing import Callable, Any, cast
from py4j.java_gateway import JavaObject

# 定义空的标记类，继承JavaObject以明确类型关系
class HadoopConfiguration(JavaObject): ...
class Path(JavaObject):
    def getFileSystem(self, conf: HadoopConfiguration) -> 'FileSystem': ...
class FileSystem(JavaObject): ...

class SparkAsset:
    def __init__(self, client):
        self._client = client
        self.path = "/some/hdfs/path"

    @property
    def hadoop_configuration(self) -> HadoopConfiguration:
        return cast(HadoopConfiguration, self._client.sparkContext._jsc.hadoopConfiguration())

    @property
    def filesystem_path_func(self) -> Callable[[str], Path]:
        return cast(Callable[[str], Path], self._client.sparkContext._jvm.org.apache.hadoop.fs.Path)

    @property
    def filesystem_path(self) -> Path:
        return self.filesystem_path_func(self.path)

    @property
    def filesystem(self) -> FileSystem:
        return cast(FileSystem, self.filesystem_path.getFileSystem(self.hadoop_configuration))

3. 封装JVM调用逻辑，隐藏类型细节 #

将操作JVM对象的代码集中到工具类，对外暴露类型明确的接口，业务代码无需直接处理JVM类型：

from typing import Any
from py4j.java_gateway import JavaObject

class HdfsUtils:
    def __init__(self, spark_context):
        self._jvm = spark_context._jvm
        self._jsc = spark_context._jsc

    def get_hadoop_config(self) -> JavaObject:
        return self._jsc.hadoopConfiguration()

    def create_path(self, path_str: str) -> JavaObject:
        return self._jvm.org.apache.hadoop.fs.Path(path_str)

    def get_filesystem(self, path: JavaObject, conf: JavaObject) -> JavaObject:
        return path.getFileSystem(conf)

    # 对外提供带明确类型的业务方法
    def path_exists(self, path_str: str) -> bool:
        conf = self.get_hadoop_config()
        path = self.create_path(path_str)
        fs = self.get_filesystem(path, conf)
        return fs.exists(path)

class SparkAsset:
    def __init__(self, client):
        self._client = client
        self.path = "/some/hdfs/path"
        self._hdfs_utils = HdfsUtils(client.sparkContext)

    @property
    def filesystem(self):
        return self._hdfs_utils.get_filesystem(
            self._hdfs_utils.create_path(self.path),
            self._hdfs_utils.get_hadoop_config()
        )

    def check_path_exists(self) -> bool:
        return self._hdfs_utils.path_exists(self.path)

4. 启用mypy的py4j插件 #

如果使用mypy做类型检查，可安装mypy-py4j插件自动识别py4j对象类型：

安装插件：

pip install mypy-py4j

在mypy.ini或pyproject.toml中配置：

[mypy]
plugins = mypy_py4j.main

该插件能自动处理部分JVM对象的类型识别，结合自定义存根使用效果更佳。

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