revise packaging check

Author: weiliangjin2021

tidy3d/packaging.py

@@ -16,7 +16,6 @@
 from tidy3d.config import config
 
 from .exceptions import Tidy3dImportError
-from .log import log
 from .version import __version__
 
 vtk = {
@@ -192,37 +191,41 @@ def _check_tidy3d_extras_available():
     Tidy3dImportError
         If tidy3d-extras is not available or not properly initialized.
     """
-    if tidy3d_extras["mod"] is None:
-        try:
-            import tidy3d_extras as tidy3d_extras_mod
+    if tidy3d_extras["mod"] is not None:
+        return
 
-        except ImportError as exc:
-            tidy3d_extras["mod"] = None
-            raise Tidy3dImportError(
-                "The package 'tidy3d-extras' is required for this "
-                "operation. Please install the 'tidy3d-extras' package using, for "
-                "example, 'pip install tidy3d[extras]'."
-            ) from exc
+    module_exists = find_spec("tidy3d_extras") is not None
+    if not module_exists:
+        raise Tidy3dImportError(
+            "The package 'tidy3d-extras' is absent. "
+            "Please install the 'tidy3d-extras' package using, for "
+            "example, 'pip install tidy3d[extras]'."
+        )
 
-        else:
-            version = tidy3d_extras_mod.__version__
+    try:
+        import tidy3d_extras as tidy3d_extras_mod
 
-            if version is None:
-                tidy3d_extras["mod"] = None
-                raise Tidy3dImportError(
-                    "The package 'tidy3d-extras' did not initialize correctly, "
-                    "likely due to an invalid API key."
-                )
+    except ImportError as exc:
+        raise Tidy3dImportError(
+            "The package 'tidy3d-extras' did not initialize correctly."
+        ) from exc
 
-            if version != __version__:
-                log.warning(
-                    "The package 'tidy3d-extras' is required for this "
-                    "operation. The version of 'tidy3d-extras' should match "
-                    "the version of 'tidy3d'. You can install the correct "
-                    "version using 'pip install tidy3d[extras]'."
-                )
+    version = tidy3d_extras_mod.__version__
+
+    if version is None:
+        raise Tidy3dImportError(
+            "The package 'tidy3d-extras' did not initialize correctly, "
+            "likely due to an invalid API key."
+        )
 
-            tidy3d_extras["mod"] = tidy3d_extras_mod
+    if version != __version__:
+        raise Tidy3dImportError(
+            f"The version of 'tidy3d-extras' is {version}, but the version of 'tidy3d' is {__version__}. "
+            "They must match. You can install the correct "
+            "version using 'pip install tidy3d[extras]'."
+        )
+
+    tidy3d_extras["mod"] = tidy3d_extras_mod
 
 
 def check_tidy3d_extras_licensed_feature(feature_name: str):
@@ -242,14 +245,16 @@ def check_tidy3d_extras_licensed_feature(feature_name: str):
     try:
         _check_tidy3d_extras_available()
     except Tidy3dImportError as exc:
-        raise Tidy3dImportError(f"Failed to load 'tidy3d-extras'. {exc!s}") from exc
-    else:
-        features = tidy3d_extras["mod"].extension._features()
-        if feature_name not in features:
-            raise Tidy3dImportError(
-                f"The feature '{feature_name}' is not available with your license. "
-                "Please contact Tidy3D support, or upgrade your license."
-            )
+        raise Tidy3dImportError(
+            f"The package 'tidy3d-extras' is required for this feature '{feature_name}'. {exc!s}"
+        ) from exc
+
+    features = tidy3d_extras["mod"].extension._features()
+    if feature_name not in features:
+        raise Tidy3dImportError(
+            f"The feature '{feature_name}' is not available with your license. "
+            "Please contact Tidy3D support, or upgrade your license."
+        )
 
 
 def supports_local_subpixel(fn):
@@ -262,18 +267,22 @@ def _fn(*args: Any, **kwargs: Any):
 
         if preference is False:
             tidy3d_extras["use_local_subpixel"] = False
-        else:
-            try:
-                _check_tidy3d_extras_available()
-            except Tidy3dImportError as exc:
-                tidy3d_extras["use_local_subpixel"] = False
-                if preference is True:
-                    raise Tidy3dImportError(
-                        f"Failed to load 'tidy3d-extras' for local subpixel support. {exc!s}"
-                    ) from exc
-            else:
-                features = tidy3d_extras["mod"].extension._features()
-                tidy3d_extras["use_local_subpixel"] = "local_subpixel" in features
+            return fn(*args, **kwargs)
+
+        try:
+            check_tidy3d_extras_licensed_feature("local_subpixel")
+        except Tidy3dImportError as exc:
+            tidy3d_extras["use_local_subpixel"] = False
+            if preference is True:
+                raise Tidy3dImportError(
+                    f"{exc!s} To suppress this error, you can set "
+                    "'config.simulation.use_local_subpixel=False'."
+                ) from exc
+            else:  # preference is None, so we can just return
+                return fn(*args, **kwargs)
+
+        # local_subpixel is available
+        tidy3d_extras["use_local_subpixel"] = True
 
         return fn(*args, **kwargs)