Adding type hints for segmentation

requirements.txt

@@ -9,3 +9,4 @@ iris
 xarray
 cartopy
 trackpy
+typing_extensions

tobac/segmentation.py

@@ -31,19 +31,28 @@
 """
 import copy
 import logging
+
+import iris.cube
 import numpy as np
+import pandas as pd
+from typing_extensions import Literal
+from typing import Union
 
 import skimage
-import numpy as np
 
 from . import utils as tb_utils
 from .utils import periodic_boundaries as pbc_utils
 from .utils import internal as internal_utils
 
 
 def add_markers(
-    features, marker_arr, seed_3D_flag, seed_3D_size=5, level=None, PBC_flag="none"
-):
+    features: pd.DataFrame,
+    marker_arr: np.array,
+    seed_3D_flag: Literal["column", "box"],
+    seed_3D_size: Union[int, tuple[int]] = 5,
+    level: Union[None, slice] = None,
+    PBC_flag: Literal["none", "hdim_1", "hdim_2", "both"] = "none",
+) -> np.array:
     """Adds markers for watershedding using the `features` dataframe
     to the marker_arr.
 
@@ -314,21 +323,21 @@ def segmentation_2D(
 
 
 def segmentation_timestep(
-    field_in,
-    features_in,
-    dxy,
-    threshold=3e-3,
-    target="maximum",
-    level=None,
-    method="watershed",
-    max_distance=None,
-    vertical_coord=None,
-    PBC_flag="none",
-    seed_3D_flag="column",
-    seed_3D_size=5,
-    segment_number_below_threshold=0,
-    segment_number_unassigned=0,
-):
+    field_in: iris.cube.Cube,
+    features_in: pd.DataFrame,
+    dxy: float,
+    threshold: float = 3e-3,
+    target: Literal["maximum", "minimum"] = "maximum",
+    level: Union[None, slice] = None,
+    method: Literal["watershed"] = "watershed",
+    max_distance: Union[None, float] = None,
+    vertical_coord: Union[str, None] = None,
+    PBC_flag: Literal["none", "hdim_1", "hdim_2", "both"] = "none",
+    seed_3D_flag: Literal["column", "box"] = "column",
+    seed_3D_size: Union[int, tuple[int]] = 5,
+    segment_number_below_threshold: int = 0,
+    segment_number_unassigned: int = 0,
+) -> tuple[iris.cube.Cube, pd.DataFrame]:
     """Perform watershedding for an individual time step of the data. Works
     for both 2D and 3D data
 
@@ -1028,7 +1037,7 @@ def check_add_unseeded_across_bdrys(
     border_max: int,
     markers_arr: np.array,
     inplace: bool = True,
-):
+) -> np.array:
     """Add new markers to unseeded but eligible regions when they are bordering
     an appropriate boundary.
 
@@ -1085,21 +1094,21 @@ def check_add_unseeded_across_bdrys(
 
 
 def segmentation(
-    features,
-    field,
-    dxy,
-    threshold=3e-3,
-    target="maximum",
-    level=None,
-    method="watershed",
-    max_distance=None,
-    vertical_coord=None,
-    PBC_flag="none",
-    seed_3D_flag="column",
-    seed_3D_size=5,
-    segment_number_below_threshold=0,
-    segment_number_unassigned=0,
-):
+    features: pd.DataFrame,
+    field: iris.cube.Cube,
+    dxy: float,
+    threshold: float = 3e-3,
+    target: Literal["maximum", "minimum"] = "maximum",
+    level: Union[None, slice] = None,
+    method: Literal["watershed"] = "watershed",
+    max_distance: Union[None, float] = None,
+    vertical_coord: Union[str, None] = None,
+    PBC_flag: Literal["none", "hdim_1", "hdim_2", "both"] = "none",
+    seed_3D_flag: Literal["column", "box"] = "column",
+    seed_3D_size: Union[int, tuple[int]] = 5,
+    segment_number_below_threshold: int = 0,
+    segment_number_unassigned: int = 0,
+) -> tuple[iris.cube.Cube, pd.DataFrame]:
     """Use watershedding to determine region above a threshold
         value around initial seeding position for all time steps of
         the input data. Works both in 2D (based on single seeding