refactor(encryption): [3/N] Move lazy import to utils/

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)
 

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

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")