add type hints to patch functions

d-chambers · d-chambers · commit bb5379ee9021 · 2025-12-18T14:39:31.000-07:00
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 = 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/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.
