Add type hints

environment-dev.yml

@@ -26,6 +26,7 @@ dependencies:
   - netcdf4
   - numpy>=1.25
   - pandas>=2.0
+  - polars>=1.0
   - pygrib>=2.1.4
   - pylint
   - pyproj>=3.6

herbie/__init__.py

@@ -42,15 +42,15 @@
     ## TODO: Will the `_version.py` file *always* be present?
     ## TODO: What if the person doesn't do "pip install"
     from ._version import __version__, __version_tuple__
-except:
+except Exception:
     __version__ = "unknown"
     __version_tuple__ = (999, 999, 999)
 
 
 ########################################################################
 # Overload Path object with my custom `expand` method so the user can
 # set environment variables in the config file (e.g., ${HOME}).
-def _expand(self, resolve=False, absolute=False):
+def _expand(self, resolve: bool = False, absolute: bool = False) -> Path:
     """
     Fully expand the Path with the given environment variables.
 
@@ -149,7 +149,7 @@ def template(self):
 try:
     # Load the Herbie config file
     config = toml.load(_config_file)
-except:
+except Exception:
     try:
         # Create the Herbie config file
         _config_path.mkdir(parents=True, exist_ok=True)

herbie/accessors.py

@@ -14,6 +14,7 @@
 import re
 import warnings
 from pathlib import Path
+from typing import Literal, Optional, Union
 
 import numpy as np
 import pandas as pd
@@ -23,7 +24,6 @@
 
 import herbie
 
-
 _level_units = dict(
     adiabaticCondensation="adiabatic condensation",
     atmosphere="atmosphere",
@@ -54,7 +54,7 @@
 )
 
 
-def add_proj_info(ds):
+def add_proj_info(ds: xr.Dataset):
     """Add projection info to a Dataset."""
     match = re.search(r'"source": "(.*?)"', ds.history)
     FILE = Path(match.group(1))
@@ -97,7 +97,7 @@ def __init__(self, xarray_obj):
         self._center = None
 
     @property
-    def center(self):
+    def center(self) -> tuple[float, float]:
         """Return the geographic center point of this dataset."""
         if self._center is None:
             # we can use a cache on our accessor objects, because accessors
@@ -107,19 +107,18 @@ def center(self):
             self._center = (float(lon.mean()), float(lat.mean()))
         return self._center
 
-    def to_180(self):
+    def to_180(self) -> xr.Dataset:
         """Wrap longitude coordinates as range [-180,180]."""
         ds = self._obj
         ds["longitude"] = (ds["longitude"] + 180) % 360 - 180
         return ds
 
-    def to_360(self):
+    def to_360(self) -> xr.Dataset:
         """Wrap longitude coordinates as range [0,360]."""
         ds = self._obj
         ds["longitude"] = (ds["longitude"] - 360) % 360
         return ds
 
-
     @functools.cached_property
     def crs(self):
         """
@@ -196,7 +195,9 @@ def polygon(self):
 
         return domain_polygon, domain_polygon_latlon
 
-    def with_wind(self, which="both"):
+    def with_wind(
+        self, which: Literal["both", "speed", "direction"] = "both"
+    ) -> xr.Dataset:
         """Return Dataset with calculated wind speed and/or direction.
 
         Consistent with the eccodes GRIB parameter database, variables
@@ -228,7 +229,7 @@ def with_wind(self, which="both"):
                 ds["si10"].attrs["standard_name"] = "wind_speed"
                 ds["si10"].attrs["grid_mapping"] = ds.u10.attrs.get("grid_mapping")
                 n_computed += 1
-            
+
             if {"u100", "v100"}.issubset(ds):
                 ds["si100"] = np.sqrt(ds.u100**2 + ds.v100**2)
                 ds["si100"].attrs["GRIB_paramId"] = 228249
@@ -237,7 +238,7 @@ def with_wind(self, which="both"):
                 ds["si100"].attrs["standard_name"] = "wind_speed"
                 ds["si100"].attrs["grid_mapping"] = ds.u100.attrs.get("grid_mapping")
                 n_computed += 1
-            
+
             if {"u80", "v80"}.issubset(ds):
                 ds["si80"] = np.sqrt(ds.u80**2 + ds.v80**2)
                 ds["si80"].attrs["long_name"] = "80 metre wind speed"
@@ -266,7 +267,7 @@ def with_wind(self, which="both"):
                 ds["wdir10"].attrs["standard_name"] = "wind_from_direction"
                 ds["wdir10"].attrs["grid_mapping"] = ds.u10.attrs.get("grid_mapping")
                 n_computed += 1
-            
+
             if {"u100", "v100"}.issubset(ds):
                 ds["wdir100"] = (
                     (270 - np.rad2deg(np.arctan2(ds.v100, ds.u100))) % 360
@@ -276,7 +277,7 @@ def with_wind(self, which="both"):
                 ds["wdir100"].attrs["standard_name"] = "wind_from_direction"
                 ds["wdir100"].attrs["grid_mapping"] = ds.u100.attrs.get("grid_mapping")
                 n_computed += 1
-            
+
             if {"u80", "v80"}.issubset(ds):
                 ds["wdir80"] = (
                     (270 - np.rad2deg(np.arctan2(ds.v80, ds.u80))) % 360
@@ -305,15 +306,15 @@ def with_wind(self, which="both"):
 
     def pick_points(
         self,
-        points,
-        method="nearest",
+        points: pd.DataFrame,
+        method: Literal["nearest", "weighted"] = "nearest",
         *,
-        k=None,
-        max_distance=500,
-        use_cached_tree=True,
-        tree_name=None,
-        verbose=False,
-    ):
+        k: Optional[int] = None,
+        max_distance: Union[int, float] = 500,
+        use_cached_tree: Union[bool, Literal["replant"]] = True,
+        tree_name: Optional[str] = None,
+        verbose: bool = False,
+    ) -> xr.Dataset:
         """Pick nearest neighbor grid values at selected  points.
 
         Parameters
@@ -384,7 +385,7 @@ def pick_points(
                 "`pip install 'herbie-data[extras]'` for the full functionality."
             )
 
-        def plant_tree(save_pickle=None):
+        def plant_tree(save_pickle: Optional[Union[Path, str]] = None):
             """Grow a new BallTree object from seedling."""
             timer = pd.Timestamp("now")
             print("INFO: 🌱 Growing new BallTree...", end="")
@@ -719,9 +720,10 @@ def plot(self, ax=None, common_features_kw={}, vars=None, **kwargs):
         raise NotImplementedError("Plotting functionality is not working right now.")
 
         try:
-            from herbie.toolbox import EasyMap, pc
-            from herbie import paint
             import matplotlib.pyplot as plt
+
+            from herbie import paint
+            from herbie.toolbox import EasyMap, pc
         except ModuleNotFoundError:
             raise ModuleNotFoundError(
                 "cartopy is an 'extra' requirement. Please use "