From 657b0fcf116cb340b44b7000221347dac70ae011 Mon Sep 17 00:00:00 2001 From: Xander Date: Mon, 14 Sep 2026 11:48:03 +0100 Subject: [PATCH] refactor(encryption): Move lazy import to utils/ --- pyiceberg/transforms.py | 33 ++++++++------------------ pyiceberg/utils/lazy_import.py | 41 +++++++++++++++++++++++++++++++++ tests/utils/test_lazy_import.py | 36 +++++++++++++++++++++++++++++ 3 files changed, 87 insertions(+), 23 deletions(-) create mode 100644 pyiceberg/utils/lazy_import.py create mode 100644 tests/utils/test_lazy_import.py diff --git a/pyiceberg/transforms.py b/pyiceberg/transforms.py index 5e0027a829..54c01d9bed 100644 --- a/pyiceberg/transforms.py +++ b/pyiceberg/transforms.py @@ -17,9 +17,7 @@ import base64 import datetime as py_datetime -import importlib import struct -import types from abc import ABC, abstractmethod from collections.abc import Callable from enum import IntEnum @@ -31,7 +29,6 @@ import mmh3 from pydantic import Field, PositiveInt, PrivateAttr -from pyiceberg.exceptions import NotInstalledError from pyiceberg.expressions import ( BoundEqualTo, BoundGreaterThan, @@ -88,6 +85,7 @@ ) from pyiceberg.utils import datetime from pyiceberg.utils.decimal import decimal_to_bytes, truncate_decimal +from pyiceberg.utils.lazy_import import try_import from pyiceberg.utils.parsing import ParseNumberFromBrackets from pyiceberg.utils.singleton import Singleton @@ -112,17 +110,6 @@ TRUNCATE_PARSER = ParseNumberFromBrackets(TRUNCATE) -def _try_import(module_name: str, extras_name: str | None = None) -> types.ModuleType: - try: - return importlib.import_module(module_name) - except ImportError: - if extras_name: - msg = f'{module_name} needs to be installed. pip install "pyiceberg[{extras_name}]"' - else: - msg = f"{module_name} needs to be installed." - raise NotInstalledError(msg) from None - - def _transform_literal(func: Callable[[Any], Any], lit: Literal[L]) -> Literal[L]: """Small helper to upwrap the value from the literal, and wrap it again.""" return literal(func(lit.value)) @@ -395,7 +382,7 @@ def __repr__(self) -> str: return f"BucketTransform(num_buckets={self._num_buckets})" def pyarrow_transform(self, source: IcebergType) -> "Callable[[pa.Array], pa.Array]": - pyiceberg_core_transform = _try_import("pyiceberg_core", extras_name="pyiceberg-core").transform + pyiceberg_core_transform = try_import("pyiceberg_core", extras_name="pyiceberg-core").transform return _pyiceberg_transform_wrapper(pyiceberg_core_transform.bucket, self._num_buckets) @@ -509,8 +496,8 @@ def __repr__(self) -> str: return "YearTransform()" def pyarrow_transform(self, source: IcebergType) -> "Callable[[pa.Array], pa.Array]": - pa = _try_import("pyarrow") - pyiceberg_core_transform = _try_import("pyiceberg_core", extras_name="pyiceberg-core").transform + pa = try_import("pyarrow") + pyiceberg_core_transform = try_import("pyiceberg_core", extras_name="pyiceberg-core").transform return _pyiceberg_transform_wrapper(pyiceberg_core_transform.year, expected_type=pa.int32()) @@ -569,8 +556,8 @@ def __repr__(self) -> str: return "MonthTransform()" def pyarrow_transform(self, source: IcebergType) -> "Callable[[pa.Array], pa.Array]": - pa = _try_import("pyarrow") - pyiceberg_core_transform = _try_import("pyiceberg_core", extras_name="pyiceberg-core").transform + pa = try_import("pyarrow") + pyiceberg_core_transform = try_import("pyiceberg_core", extras_name="pyiceberg-core").transform return _pyiceberg_transform_wrapper(pyiceberg_core_transform.month, expected_type=pa.int32()) @@ -638,8 +625,8 @@ def __repr__(self) -> str: return "DayTransform()" def pyarrow_transform(self, source: IcebergType) -> "Callable[[pa.Array], pa.Array]": - pa = _try_import("pyarrow", extras_name="pyarrow") - pyiceberg_core_transform = _try_import("pyiceberg_core", extras_name="pyiceberg-core").transform + pa = try_import("pyarrow", extras_name="pyarrow") + pyiceberg_core_transform = try_import("pyiceberg_core", extras_name="pyiceberg-core").transform return _pyiceberg_transform_wrapper(pyiceberg_core_transform.day, expected_type=pa.int32()) @@ -691,7 +678,7 @@ def __repr__(self) -> str: return "HourTransform()" def pyarrow_transform(self, source: IcebergType) -> "Callable[[pa.Array], pa.Array]": - pyiceberg_core_transform = _try_import("pyiceberg_core", extras_name="pyiceberg-core").transform + pyiceberg_core_transform = try_import("pyiceberg_core", extras_name="pyiceberg-core").transform return _pyiceberg_transform_wrapper(pyiceberg_core_transform.hour) @@ -918,7 +905,7 @@ def __repr__(self) -> str: return f"TruncateTransform(width={self._width})" def pyarrow_transform(self, source: IcebergType) -> "Callable[[pa.Array], pa.Array]": - pyiceberg_core_transform = _try_import("pyiceberg_core", extras_name="pyiceberg-core").transform + pyiceberg_core_transform = try_import("pyiceberg_core", extras_name="pyiceberg-core").transform return _pyiceberg_transform_wrapper(pyiceberg_core_transform.truncate, self._width) diff --git a/pyiceberg/utils/lazy_import.py b/pyiceberg/utils/lazy_import.py new file mode 100644 index 0000000000..8f6d7496c0 --- /dev/null +++ b/pyiceberg/utils/lazy_import.py @@ -0,0 +1,41 @@ +# 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. +"""Helpers for importing modules that ship in an optional extra.""" + +from __future__ import annotations + +import importlib +import types + +from pyiceberg.exceptions import NotInstalledError + + +def try_import(module_name: str, extras_name: str | None = None) -> types.ModuleType: + """Import `module_name`, raising `NotInstalledError` with an install hint when it is missing. + + Args: + module_name (str): The module to import. + extras_name (str | None): The pyiceberg extra that provides it, if any. + """ + try: + return importlib.import_module(module_name) + except ImportError: + if extras_name: + msg = f'{module_name} needs to be installed. pip install "pyiceberg[{extras_name}]"' + else: + msg = f"{module_name} needs to be installed." + raise NotInstalledError(msg) from None diff --git a/tests/utils/test_lazy_import.py b/tests/utils/test_lazy_import.py new file mode 100644 index 0000000000..d68da5423f --- /dev/null +++ b/tests/utils/test_lazy_import.py @@ -0,0 +1,36 @@ +# 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 importlib + +import pytest + +from pyiceberg.exceptions import NotInstalledError +from pyiceberg.utils.lazy_import import try_import + + +def test_try_import_returns_the_module() -> None: + assert try_import("importlib") is importlib + + +def test_try_import_missing_module_with_extras() -> None: + with pytest.raises(NotInstalledError, match=r'nonexistent needs to be installed. pip install "pyiceberg\[some-extra\]"'): + try_import("nonexistent", extras_name="some-extra") + + +def test_try_import_missing_module_without_extras() -> None: + with pytest.raises(NotInstalledError, match="nonexistent needs to be installed."): + try_import("nonexistent")