Move optional imports inside the functions and classes and raise inst…

…ead of warn if the libraries aren't available
blaylockbk · May 28, 2024 · fecf307 · fecf307
1 parent a9c835d
commit fecf307
Show file tree

Hide file tree

Showing 3 changed files with 41 additions and 39 deletions.
diff --git a/herbie/accessors.py b/herbie/accessors.py
@@ -23,32 +23,6 @@
 
 import herbie
 
-# Optional dependencies
-try:
-    import metpy
-    import matplotlib.pyplot as plt
-except ModuleNotFoundError:
-    warnings.warn(
-        "metpy is an 'extra' requirement, please use "
-        "`pip install 'herbie-data[extras]'` for the full functionality."
-)
-try:
-    import cartopy.crs as ccrs
-    import shapely
-    from shapely.geometry import MultiPoint, Point, Polygon
-except ModuleNotFoundError:
-    warnings.warn(
-        "cartopy is an 'extra' requirements, please use "
-        "`pip install 'herbie-data[extras]'` for the full functionality."
-    )
-try:
-    from sklearn.neighbors import BallTree
-except ModuleNotFoundError:
-    warnings.warn(
-        "scikit-learn is an 'extra' requirement, please use "
-        "`pip install 'herbie-data[extras]'` for the full functionality."
-    )
-
 
 _level_units = dict(
     adiabaticCondensation="adiabatic condensation",
@@ -121,6 +95,24 @@ class HerbieAccessor:
     def __init__(self, xarray_obj):
         self._obj = xarray_obj
         self._center = None
+        # Optional dependencies
+        try:
+            import metpy
+            import matplotlib.pyplot as plt
+        except ModuleNotFoundError:
+            raise ModuleNotFoundError(
+                "metpy is an 'extra' requirement, please use "
+                "`pip install 'herbie-data[extras]'` for the full functionality."
+        )
+        try:
+            import cartopy.crs as ccrs
+            import shapely
+            from shapely.geometry import MultiPoint, Point, Polygon
+        except ModuleNotFoundError:
+            raise ModuleNotFoundError(
+                "cartopy is an 'extra' requirements, please use "
+                "`pip install 'herbie-data[extras]'` for the full functionality."
+            )
 
     @property
     def center(self):
@@ -268,6 +260,13 @@ def pick_points(
 
         def plant_tree(save_pickle=None):
             """Grow a new BallTree object from seedling."""
+            try:
+                from sklearn.neighbors import BallTree
+            except ModuleNotFoundError:
+                raise ModuleNotFoundError(
+                    "scikit-learn is an 'extra' requirement, please use "
+                    "`pip install 'herbie-data[extras]'` for the full functionality."
+                )
             timer = pd.Timestamp("now")
             print("INFO: 🌱 Growing new BallTree...", end="")
             tree = BallTree(np.deg2rad(df_grid), metric="haversine")

diff --git a/herbie/hrrr_zarr.py b/herbie/hrrr_zarr.py
@@ -12,7 +12,6 @@
 """
 
 import datetime
-import warnings
 import pandas as pd
 import s3fs
 import xarray as xr
@@ -28,7 +27,7 @@ def load_dataset(urls):
     try:
         import cartopy.crs as ccrs
     except ModuleNotFoundError:
-        warnings.warn(
+        raise ModuleNotFoundError(
             "cartopy is an 'extra' requirement, please use "
             "`pip install 'herbie-data[extras]'` for the full functionality."
     )

diff --git a/herbie/misc.py b/herbie/misc.py
@@ -18,17 +18,6 @@
   ██  ██  ██  ██  ██  ██  ██  ██  ██  ██  ██  ██  ██  ██  ██  ██  ██  ██
 """
 
-import warnings
-
-try:
-    import matplotlib.pyplot as plt
-    import matplotlib.patheffects as path_effects
-except ModuleNotFoundError:
-    warnings.warn(
-        "matplotlib is an 'extra' requirement, please use "
-        "`pip install 'herbie-data[extras]'` for the full functionality."
-)
-
 
 class hc:
     """Herbie Color Pallette"""
@@ -130,6 +119,13 @@ def print_rich(H):
 
 def HerbieLogo(white_line=False):
     """Logo of Herbie The Love Bug"""
+    try:
+        import matplotlib.pyplot as plt
+    except ModuleNotFoundError:
+        raise ModuleNotFoundError(
+            "matplotlib is an 'extra' requirement, please use "
+            "`pip install 'herbie-data[extras]'` for the full functionality."
+    )
     plt.figure(figsize=[5, 5], facecolor=hc.tan)
 
     plt.axis([-10, 10, -10, 10])
@@ -179,6 +175,14 @@ def HerbieLogo2(white_line=False, text_color="tan", text_stroke="black"):
     >>> ax = HerbieLogo2(text_color='tan')
     >>> plt.savefig('Herbie_transparent_tan.svg', bbox_inches="tight", transparent=True)
     """
+    try:
+        import matplotlib.pyplot as plt
+        import matplotlib.patheffects as path_effects
+    except ModuleNotFoundError:
+        raise ModuleNotFoundError(
+            "matplotlib is an 'extra' requirement, please use "
+            "`pip install 'herbie-data[extras]'` for the full functionality."
+    )
     plt.figure(figsize=[5, 3], facecolor=hc.tan)
 
     plt.axis([1.5, 20, -10, 10])