_importing.py

# Copyright 2025 Google LLC
#
# Licensed 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 importlib
from types import ModuleType

import numpy
from packaging import version

# Keep this in sync with setup.py
POLARS_MIN_VERSION = version.Version("1.7.0")


def import_polars() -> ModuleType:
    polars_module = importlib.import_module("polars")
    # Check for necessary methods instead of the version number because we
    # can't trust the polars version until
    # https://github.com/pola-rs/polars/issues/23940 is fixed.
    try:
        polars_module.lit(numpy.int64(100), dtype=polars_module.Int64())
    except TypeError:
        raise ImportError(
            f"Imported polars version is likely below the minimum version: {POLARS_MIN_VERSION}"
        )
    return polars_module


DATAFUSION_MIN_VERSION = version.Version("52.0.0")


def import_datafusion() -> ModuleType:
    datafusion_module = importlib.import_module("datafusion")
    # Add any version checks if necessary, for now just check it imports
    df_version = version.Version(datafusion_module.__version__)
    if df_version < DATAFUSION_MIN_VERSION:
        raise ImportError(
            f"Imported datafusion version {df_version} is below the minimum version: {DATAFUSION_MIN_VERSION}"
        )
    return datafusion_module