add type hints to patch functions, fix febus utils and trapz (#590)

d-chambers · web-flow · commit ebdbecb09c9e · 2025-12-22T10:16:22.000-07:00
* add type hints to patch functions

* review

* debug profile

* try fix benchmarks

* continue debug febus block time

* more unpacking

* fix trapz issue

* remove trapz reference

* global trap function
diff --git a/.github/workflows/profile.yml b/.github/workflows/profile.yml
@@ -33,4 +33,4 @@ jobs:
         with:
           mode: instrumentation
           run: ./.github/test_code.sh profile
-          token: ${{ secrets.CODSPEED_TOKEN }}  # Optional for public repos
+          token: ${{ secrets.CODSPEED_TOKEN }}
diff --git a/dascore/io/febus/utils.py b/dascore/io/febus/utils.py
@@ -33,9 +33,9 @@ def _get_block_time(feb):
     # Some files have this set. We haven't yet seen any files where this
     # values exists and is wrong, so we trust it (for now). This is probably
     # much faster than reading the whole time vector.
-    br = feb.zone.attrs.get("BlockRate", 0) / 1_000
+    br = _maybe_unpack(feb.zone.attrs.get("BlockRate", 0) / 1_000)
     if br > 0:
-        return 1 / br
+        return float(1 / br)
     # Otherwise we have to try to use the time vector. Here be dragons.
     time_shape = feb.source["time"].shape
     # Not sure why but time has the shape of [1, n] for some files and just
@@ -54,7 +54,7 @@ def _get_block_time(feb):
     # After removing outliers, the mean seems to work better than the median
     # for the test files we have. There is still a concerning amount of
     # variability.
-    return np.mean(d_time)
+    return float(_maybe_unpack(np.mean(d_time)))
 
 
 @cache
diff --git a/dascore/proc/coords.py b/dascore/proc/coords.py
@@ -22,7 +22,7 @@
 
 
 @patch_function()
-def snap_coords(patch: PatchType, *coords, reverse=False):
+def snap_coords(patch: PatchType, *coords, reverse: bool = False) -> PatchType:
     """
     Snap coordinates to evenly spaced samples.
 
@@ -59,7 +59,7 @@ def snap_coords(patch: PatchType, *coords, reverse=False):
 
 
 @patch_function()
-def sort_coords(patch: PatchType, *coords, reverse=False):
+def sort_coords(patch: PatchType, *coords, reverse: bool = False) -> PatchType:
     """
     Sort one or more coordinates.
 
diff --git a/dascore/proc/correlate.py b/dascore/proc/correlate.py
@@ -35,7 +35,9 @@ def _get_source_fft(patch, dim, source, source_axis, samples):
 
 
 @patch_function()
-def correlate_shift(patch, dim, undo_weighting=True):
+def correlate_shift(
+    patch: PatchType, dim: str, undo_weighting: bool = True
+) -> PatchType:
     """
     Apply a shift to the patch data to undo correlation in frequency domain.
 
@@ -87,7 +89,7 @@ def correlate_shift(patch, dim, undo_weighting=True):
 @patch_function()
 def correlate(
     patch: PatchType,
-    samples=False,
+    samples: bool = False,
     lag=None,
     **kwargs,
 ) -> PatchType:
diff --git a/dascore/proc/detrend.py b/dascore/proc/detrend.py
@@ -2,14 +2,18 @@
 
 from __future__ import annotations
 
+from typing import Literal
+
 from scipy.signal import detrend as scipy_detrend
 
 from dascore.constants import PatchType
 from dascore.utils.patch import patch_function
 
 
 @patch_function()
-def detrend(patch: PatchType, dim, type="linear") -> PatchType:
+def detrend(
+    patch: PatchType, dim: str, type: Literal["linear", "constant"] = "linear"
+) -> PatchType:
     """
     Perform detrending along a given dimension (distance or time) of a patch.
 
diff --git a/dascore/proc/filter.py b/dascore/proc/filter.py
@@ -95,7 +95,9 @@ def _get_sos(sr, filt_min, filt_max, corners):
 
 
 @patch_function()
-def pass_filter(patch: PatchType, corners=4, zerophase=True, **kwargs) -> PatchType:
+def pass_filter(
+    patch: PatchType, corners: int = 4, zerophase: bool = True, **kwargs
+) -> PatchType:
     """
     Apply a Butterworth pass filter (bandpass, highpass, or lowpass).
 
@@ -146,7 +148,9 @@ def pass_filter(patch: PatchType, corners=4, zerophase=True, **kwargs) -> PatchT
 
 
 @patch_function()
-def sobel_filter(patch: PatchType, dim: str, mode="reflect", cval=0.0) -> PatchType:
+def sobel_filter(
+    patch: PatchType, dim: str, mode: str = "reflect", cval: float | int = 0.0
+) -> PatchType:
     """
     Apply a Sobel filter.
 
@@ -200,7 +204,11 @@ def _create_size_and_axes(patch, kwargs, samples):
 @patch_function()
 @compose_docstring(sample_explanation=samples_arg_description)
 def median_filter(
-    patch: PatchType, samples=False, mode="reflect", cval=0.0, **kwargs
+    patch: PatchType,
+    samples: bool = False,
+    mode: str = "reflect",
+    cval: float = 0.0,
+    **kwargs,
 ) -> PatchType:
     """
     Apply 2-D median filter.
@@ -251,7 +259,7 @@ def median_filter(
 
 
 @patch_function()
-def notch_filter(patch: PatchType, q, **kwargs) -> PatchType:
+def notch_filter(patch: PatchType, q: float, **kwargs) -> PatchType:
     """
     Apply a second-order IIR notch digital filter on patch's data.
 
@@ -320,7 +328,12 @@ def notch_filter(patch: PatchType, q, **kwargs) -> PatchType:
 @patch_function()
 @compose_docstring(sample_explanation=samples_arg_description)
 def savgol_filter(
-    patch: PatchType, polyorder, samples=False, mode="interp", cval=0.0, **kwargs
+    patch: PatchType,
+    polyorder: int,
+    samples: bool = False,
+    mode: str = "interp",
+    cval: float = 0.0,
+    **kwargs,
 ) -> PatchType:
     """
     Applies Savgol filter along spenfied dimensions.
@@ -382,7 +395,12 @@ def savgol_filter(
 @patch_function()
 @compose_docstring(sample_explanation=samples_arg_description)
 def gaussian_filter(
-    patch: PatchType, samples=False, mode="reflect", cval=0.0, truncate=4.0, **kwargs
+    patch: PatchType,
+    samples: bool = False,
+    mode: str = "reflect",
+    cval: float = 0.0,
+    truncate: float = 4.0,
+    **kwargs,
 ) -> PatchType:
     """
     Applies a Gaussian filter along specified dimensions.
diff --git a/dascore/proc/hampel.py b/dascore/proc/hampel.py
@@ -99,10 +99,10 @@ def hampel_filter(
     patch: PatchType,
     *,
     threshold: float = 10.0,
-    samples=False,
-    approximate=True,
+    samples: bool = False,
+    approximate: bool = True,
     **kwargs,
-):
+) -> PatchType:
     """
     A Hampel filter implementation useful for removing spikes in data.
 
diff --git a/dascore/proc/resample.py b/dascore/proc/resample.py
@@ -40,7 +40,7 @@ def _apply_scipy_decimation(patch, factor, ftype, axis):
 def decimate(
     patch: PatchType,
     filter_type: Literal["iir", "fir", None] = "iir",
-    copy=True,
+    copy: bool = True,
     **kwargs,
 ) -> PatchType:
     """
@@ -155,7 +155,11 @@ def interpolate(patch: PatchType, kind: str | int = "linear", **kwargs) -> Patch
 
 @patch_function()
 def resample(
-    patch: PatchType, window=None, interp_kind="linear", samples=False, **kwargs
+    patch: PatchType,
+    window=None,
+    interp_kind: str = "linear",
+    samples: bool = False,
+    **kwargs,
 ) -> PatchType:
     """
     Resample along a single dimension using Fourier Method and interpolation.
diff --git a/dascore/proc/wiener.py b/dascore/proc/wiener.py
@@ -15,10 +15,10 @@
 def wiener_filter(
     patch: PatchType,
     *,
-    noise=None,
-    samples=False,
+    noise: float | None = None,
+    samples: bool = False,
     **kwargs,
-):
+) -> PatchType:
     """
     Apply a Wiener filter to reduce noise in the patch data.
 
diff --git a/dascore/transform/integrate.py b/dascore/transform/integrate.py
@@ -16,6 +16,8 @@
     patch_function,
 )
 
+_TRAP_FUNC = getattr(np, "trapezoid" if hasattr(np, "trapezoid") else "trapz")
+
 
 def _quasi_mean(array):
     """Get a quasi mean value from an array. Works with datetimes."""
@@ -45,12 +47,11 @@ def _get_new_coords_and_array(patch, array, dims):
     ndims = len(patch.shape)
     for dxs_or_val, ax in zip(dxs_or_vals, axes):
         # Numpy 2/3 compat code
-        trap = getattr(np, "trapezoid", getattr(np, "trapz"))
         indexer = broadcast_for_index(ndims, ax, None, fill=slice(None))
         if is_array(dxs_or_val):
-            array = trap(array, x=dxs_or_val, axis=ax)[indexer]
+            array = _TRAP_FUNC(array, x=dxs_or_val, axis=ax)[indexer]
         else:
-            array = trap(array, dx=dxs_or_val, axis=ax)[indexer]
+            array = _TRAP_FUNC(array, dx=dxs_or_val, axis=ax)[indexer]
     array, coords = _get_new_coords_and_array(patch, array, dims)
     return array, coords
 
diff --git a/dascore/utils/misc.py b/dascore/utils/misc.py
@@ -461,9 +461,10 @@ def maybe_get_items(
 
 
 def _maybe_unpack(maybe_array):
-    """Unpack an array like object if it is size one, else return input."""
+    """Unpack a single-element array-like object, else return input unchanged."""
     size = getattr(maybe_array, "size", 0)
-    if size == 1:
+    shape = getattr(maybe_array, "shape", ())
+    if size == 1 and shape:
         maybe_array = maybe_array[0]
     return maybe_array
 
diff --git a/dascore/viz/waterfall.py b/dascore/viz/waterfall.py
@@ -135,7 +135,7 @@ def _add_colorbar(ax, im, patch, log):
 def waterfall(
     patch: PatchType,
     ax: plt.Axes | None = None,
-    cmap="bwr",
+    cmap: str = "bwr",
     scale: float | Sequence[float] | None = None,
     scale_type: Literal["relative", "absolute"] = "relative",
     interpolation: str | None = "antialiased",
diff --git a/dascore/viz/wiggle.py b/dascore/viz/wiggle.py
@@ -53,13 +53,13 @@ def _format_y_axis_ticks(ax, offsets, other_axis_ticks, max_ticks=10):
 @patch_function()
 def wiggle(
     patch: PatchType,
-    dim="time",
-    scale=1,
-    alpha=0.2,
-    color="black",
-    shade=False,
+    dim: str = "time",
+    scale: float = 1,
+    alpha: float = 0.2,
+    color: str = "black",
+    shade: bool = False,
     ax: plt.Axes | None = None,
-    show=False,
+    show: bool = False,
 ) -> plt.Figure:
     """
     Create a wiggle plot of patch data.
diff --git a/tests/test_transform/test_integrate.py b/tests/test_transform/test_integrate.py
@@ -91,7 +91,8 @@ def test_simple_integration(self, ones_patch):
             out = patch.integrate(dim=dim, definite=True)
             assert out.shape[ax] == 1
             step = to_float(patch.get_coord(dim).step)
-            trap = getattr(np, "trapezoid", getattr(np, "trapz"))
+            trap_name = "trapezoid" if hasattr(np, "trapezoid") else "trapz"
+            trap = getattr(np, trap_name)
             expected_data = trap(patch.data, dx=step, axis=ax)
             ndims = len(patch.dims)
             indexer = broadcast_for_index(ndims, ax, None)
