refactor: centralize SNP data type aliases in base_params.py

Author: Yashsingh045

malariagen_data/anoph/base_params.py

@@ -10,6 +10,7 @@
     single_contig_param_type,
     single_region_param_type,
     chunks_param_type,
+    Region,
 )
 
 contig: TypeAlias = Annotated[
@@ -47,6 +48,8 @@
     """,
 ]
 
+regions_tuple: TypeAlias = Tuple[Region, ...]
+
 release: TypeAlias = Annotated[
     Union[str, Sequence[str]],
     "Release version identifier.",
@@ -65,6 +68,8 @@
     """,
 ]
 
+sample_sets_tuple: TypeAlias = Tuple[sample_set, ...]
+
 sample_query: TypeAlias = Annotated[
     str,
     """
@@ -94,6 +99,8 @@
     """,
 ]
 
+sample_indices_tuple: TypeAlias = Tuple[int, ...]
+
 sample: TypeAlias = Annotated[
     Union[str, int],
     "Sample identifier or index within sample set.",

malariagen_data/anoph/snp_data.py

@@ -111,7 +111,7 @@ def _prep_site_mask_param(
         self,
         *,
         site_mask: base_params.site_mask,
-    ) -> str:
+    ) -> base_params.site_mask:
         if site_mask == base_params.DEFAULT:
             # Use whatever is the default site mask for this data resource.
             assert self._default_site_mask is not None
@@ -125,7 +125,7 @@ def _prep_optional_site_mask_param(
         self,
         *,
         site_mask: Optional[base_params.site_mask],
-    ) -> Optional[str]:
+    ) -> Optional[base_params.site_mask]:
         if site_mask is None:
             # This is allowed, it means don't apply any site mask to the data.
             return None
@@ -1113,9 +1113,9 @@ def snp_calls(
     def _raw_snp_calls(
         self,
         *,
-        regions: Tuple[Region, ...],
-        sample_sets: Optional[Tuple[str, ...]],
-        site_mask: Optional[str],
+        regions: base_params.regions_tuple,
+        sample_sets: Optional[base_params.sample_sets_tuple],
+        site_mask: Optional[base_params.site_mask],
         site_class: Optional[base_params.site_class],
         inline_array: base_params.inline_array,
         chunks: base_params.chunks,
@@ -1182,10 +1182,10 @@ def _raw_snp_calls(
     def _snp_calls(
         self,
         *,
-        regions: Tuple[Region, ...],
-        sample_sets: Optional[Tuple[str, ...]],
-        sample_indices: Optional[Tuple[int, ...]],
-        site_mask: Optional[str],
+        regions: base_params.regions_tuple,
+        sample_sets: Optional[base_params.sample_sets_tuple],
+        sample_indices: Optional[base_params.sample_indices_tuple],
+        site_mask: Optional[base_params.site_mask],
         site_class: Optional[base_params.site_class],
         cohort_size: Optional[base_params.cohort_size],
         min_cohort_size: Optional[base_params.min_cohort_size],
@@ -1309,9 +1309,9 @@ def _snp_allele_counts(
         self,
         *,
         region: Union[dict, List[dict]],
-        sample_sets: Optional[Tuple[str, ...]],
-        sample_indices: Optional[Tuple[int, ...]],
-        site_mask: Optional[str],
+        sample_sets: Optional[base_params.sample_sets_tuple],
+        sample_indices: Optional[base_params.sample_indices_tuple],
+        site_mask: Optional[base_params.site_mask],
         site_class: Optional[base_params.site_class],
         cohort_size: Optional[base_params.cohort_size],
         min_cohort_size: Optional[base_params.min_cohort_size],