python/pyspark/sql/connect/utils.py - master - githubmirror/spark

python/pyspark/sql/connect/utils.py
123 строки4 KB
Tian Gao
[SPARK-55667][PYTHON][CONNECT][FOLLOWUP] Detect REPL for doctest check
20 мар 2026, 02:25
20 мар 2026, 02:25a3eff22
#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
import sys

from pyspark.loose_version import LooseVersion
from pyspark.sql.pandas.utils import require_minimum_pandas_version, require_minimum_pyarrow_version
from pyspark.errors import PySparkImportError


def check_dependencies() -> None:
    main_module = sys.modules["__main__"]
    if (
        main_module.__spec__ is None
        and not hasattr(main_module, "__file__")
        and not hasattr(sys, "ps1")
    ):
        # The main module is not initialized at all at this point. We must be running doctests.
        from pyspark.testing.connectutils import should_test_connect, connect_requirement_message

        if not should_test_connect:
            print(
                f"Skipping doctests: {connect_requirement_message}",
                file=sys.stderr,
            )
            sys.exit(0)

    require_minimum_pandas_version()
    require_minimum_pyarrow_version()
    require_minimum_grpc_version()
    require_minimum_grpcio_status_version()
    require_minimum_googleapis_common_protos_version()
    require_minimum_zstandard_version()


def require_minimum_grpc_version() -> None:
    """Raise ImportError if minimum version of grpc is not installed"""
    minimum_grpc_version = "1.48.1"

    try:
        import grpc
    except ImportError as error:
        raise PySparkImportError(
            errorClass="PACKAGE_NOT_INSTALLED",
            messageParameters={
                "package_name": "grpcio",
                "minimum_version": str(minimum_grpc_version),
            },
        ) from error
    if LooseVersion(grpc.__version__) < LooseVersion(minimum_grpc_version):
        raise PySparkImportError(
            errorClass="UNSUPPORTED_PACKAGE_VERSION",
            messageParameters={
                "package_name": "grpcio",
                "minimum_version": str(minimum_grpc_version),
                "current_version": str(grpc.__version__),
            },
        )


def require_minimum_grpcio_status_version() -> None:
    """Raise ImportError if grpcio-status is not installed"""
    minimum_grpc_version = "1.48.1"

    try:
        import grpc_status  # noqa
    except ImportError as error:
        raise PySparkImportError(
            errorClass="PACKAGE_NOT_INSTALLED",
            messageParameters={
                "package_name": "grpcio-status",
                "minimum_version": str(minimum_grpc_version),
            },
        ) from error


def require_minimum_googleapis_common_protos_version() -> None:
    """Raise ImportError if googleapis-common-protos is not installed"""
    minimum_common_protos_version = "1.56.4"

    try:
        import google.rpc  # noqa
    except ImportError as error:
        raise PySparkImportError(
            errorClass="PACKAGE_NOT_INSTALLED",
            messageParameters={
                "package_name": "googleapis-common-protos",
                "minimum_version": str(minimum_common_protos_version),
            },
        ) from error


def require_minimum_zstandard_version() -> None:
    """Raise ImportError if zstandard is not installed"""
    minimum_zstandard_version = "0.25.0"

    try:
        import zstandard  # noqa
    except ImportError as error:
        raise PySparkImportError(
            errorClass="PACKAGE_NOT_INSTALLED",
            messageParameters={
                "package_name": "zstandard",
                "minimum_version": str(minimum_zstandard_version),
            },
        ) from error


def get_python_ver() -> str:
    return "%d.%d" % sys.version_info[:2]