(iceberg-python) branch main updated: refactor(encryption): [3/N] Move lazy import to utils/ (#3961)

[gi...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg-python.git


The following commit(s) were added to refs/heads/main by this push:
     new a2cbb8252 refactor(encryption): [3/N] Move lazy import to utils/ (#3961)
a2cbb8252 is described below

commit a2cbb8252df913318b217bcf906c56898b9f4b55
Author: Xander <za...@googlemail.com>
AuthorDate: Mon Sep 14 15:35:44 2026 +0000

    refactor(encryption): [3/N] Move lazy import to utils/ (#3961)
    
    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(-)

diff --git a/pyiceberg/transforms.py b/pyiceberg/transforms.py
index 5e0027a82..54c01d9be 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 @@ from uuid import UUID
 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.types import (
 )
 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 @@ BUCKET_PARSER = ParseNumberFromBrackets(BUCKET)
 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 @@ class BucketTransform(Transform[S, int]):
         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 @@ class YearTransform(TimeTransform[S]):
         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 @@ class MonthTransform(TimeTransform[S]):
         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 @@ class DayTransform(TimeTransform[S]):
         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 @@ class HourTransform(TimeTransform[S]):
         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 @@ class TruncateTransform(Transform[S, S]):
         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 000000000..8f6d7496c
--- /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 000000000..d68da5423
--- /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")