Declare package typed

bopen · Apr 25, 2022 · 5fa8bb4 · 5fa8bb4
1 parent 38772b7
commit 5fa8bb4
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 6 deletions.
diff --git a/tests/test_20_sentinel1.py b/tests/test_20_sentinel1.py
@@ -295,12 +295,16 @@ def test_crop_burst_dataset() -> None:
 
 
 def test_mosaic_slc_iw() -> None:
-    da = sentinel1.open_sentinel1_dataset(SLC_IW_V340, group="IW1/HH")
+    ds = sentinel1.open_sentinel1_dataset(SLC_IW_V340, group="IW1/HH")
 
-    res = sentinel1.mosaic_slc_iw(da)
+    res = sentinel1.mosaic_slc_iw(ds)
 
     assert isinstance(res, xr.Dataset)
 
+    res = sentinel1.mosaic_slc_iw(ds.measurement)
+
+    assert isinstance(res, xr.DataArray)
+
 
 def test_calibrate_amplitude() -> None:
     swath_ds = sentinel1.open_sentinel1_dataset(SLC_IW, group="IW1/VH")

diff --git a/xarray_sentinel/py.typed b/xarray_sentinel/py.typed
diff --git a/xarray_sentinel/sentinel1.py b/xarray_sentinel/sentinel1.py
@@ -24,6 +24,9 @@
 ONE_SECOND = np.timedelta64(1, "s")
 
 
+DataArrayOrDataset = T.TypeVar("DataArrayOrDataset", xr.DataArray, xr.Dataset)
+
+
 def get_fs_path(
     urlpath_or_path: esa_safe.PathType,
     fs: T.Optional[fsspec.AbstractFileSystem] = None,
@@ -551,7 +554,7 @@ def open_pol_dataset(
 
 
 def find_bursts_index(
-    pol_dataset: xr.Dataset,
+    pol_dataset: DataArrayOrDataset,
     azimuth_anx_time: float,
     use_center: bool = False,
 ) -> int:
@@ -573,12 +576,12 @@ def find_bursts_index(
 
 
 def crop_burst_dataset(
-    pol_dataset: xr.Dataset,
+    pol_dataset: DataArrayOrDataset,
     burst_index: T.Optional[int] = None,
     azimuth_anx_time: T.Optional[float] = None,
     use_center: bool = False,
     burst_id: T.Optional[int] = None,
-) -> xr.Dataset:
+) -> DataArrayOrDataset:
     burst_definitions = (
         (burst_index is not None)
         + (azimuth_anx_time is not None)
@@ -632,7 +635,9 @@ def crop_burst_dataset(
     return ds
 
 
-def mosaic_slc_iw(slc_iw_image: xr.Dataset, crop: int = 90) -> xr.Dataset:
+def mosaic_slc_iw(
+    slc_iw_image: DataArrayOrDataset, crop: int = 90
+) -> DataArrayOrDataset:
     bursts = []
     for i in range(slc_iw_image.attrs["number_of_bursts"]):
         burst = crop_burst_dataset(slc_iw_image, burst_index=i)