handle set of dataset types during allele type validation

Author: bpblanken

v03_pipeline/lib/misc/validation.py

@@ -1,3 +1,4 @@
+from collections.abc import Iterable
 from typing import Any
 
 import hail as hl
@@ -27,25 +28,29 @@ def __init__(
 
 def validate_allele_type(
     t: hl.Table | hl.MatrixTable,
-    dataset_type: DatasetType,
+    dataset_type: DatasetType | Iterable[DatasetType],
     **_: Any,
 ) -> None:
     ht = t.rows() if isinstance(t, hl.MatrixTable) else t
-    ht = ht.filter(
-        dataset_type.invalid_allele_types.contains(
-            hl.numeric_allele_type(ht.alleles[0], ht.alleles[1]),
-        ),
+    dataset_types = (
+        [dataset_type] if isinstance(dataset_type, DatasetType) else dataset_type
     )
-    if ht.count() > 0:
-        collected_alleles = sorted(
-            [tuple(x) for x in ht.aggregate(hl.agg.collect_as_set(ht.alleles))],
+    for dataset_type in dataset_types:
+        ht = ht.filter(
+            dataset_type.invalid_allele_types.contains(
+                hl.numeric_allele_type(ht.alleles[0], ht.alleles[1]),
+            ),
         )
-        # Handle case where all invalid alleles are NON_REF, indicating a gvcf:
-        if all('<NON_REF>' in alleles for alleles in collected_alleles):
-            msg = 'Alleles with invalid allele <NON_REF> are present in the callset.  This appears to be a GVCF containing records for sites with no variants.'
+        if ht.count() > 0:
+            collected_alleles = sorted(
+                [tuple(x) for x in ht.aggregate(hl.agg.collect_as_set(ht.alleles))],
+            )
+            # Handle case where all invalid alleles are NON_REF, indicating a gvcf:
+            if all('<NON_REF>' in alleles for alleles in collected_alleles):
+                msg = 'Alleles with invalid allele <NON_REF> are present in the callset.  This appears to be a GVCF containing records for sites with no variants.'
+                raise SeqrValidationError(msg)
+            msg = f'Alleles with invalid AlleleType are present in the callset: {collected_alleles[:10]}'
             raise SeqrValidationError(msg)
-        msg = f'Alleles with invalid AlleleType are present in the callset: {collected_alleles[:10]}'
-        raise SeqrValidationError(msg)
 
 
 def validate_no_duplicate_variants(

v03_pipeline/lib/misc/validation_test.py

@@ -85,7 +85,7 @@ def test_validate_allele_type(self) -> None:
             "Alleles with invalid AlleleType are present in the callset: \\[\\('A', '-'\\), \\('A', '<NON_REF>'\\)\\]",
             validate_allele_type,
             mt,
-            DatasetType.SNV_INDEL,
+            [DatasetType.SNV_INDEL],
         )
 
         mt = (
@@ -119,7 +119,7 @@ def test_validate_allele_type(self) -> None:
             'Alleles with invalid allele <NON_REF> are present in the callset.  This appears to be a GVCF containing records for sites with no variants.',
             validate_allele_type,
             mt,
-            DatasetType.SNV_INDEL,
+            [DatasetType.SNV_INDEL],
         )
 
     def test_validate_imputed_sex_ploidy(self) -> None:

v03_pipeline/lib/reference_datasets/reference_dataset.py

@@ -104,6 +104,12 @@ def version(self, reference_genome: ReferenceGenome) -> str:
             )
         return version
 
+    def dataset_types(
+        self,
+        reference_genome: ReferenceGenome,
+    ) -> frozenset[DatasetType]:
+        return CONFIG[self][reference_genome][DATASET_TYPES]
+
     @property
     def enums(self) -> dict | None:
         return CONFIG[self].get(ENUMS)
@@ -143,10 +149,7 @@ def get_ht(
         if enum_selects:
             ht = ht.transmute(**enum_selects)
         ht = filter_contigs(ht, reference_genome)
-        # Reference Datasets are DatasetType agnostic, but these
-        # methods (in theory) support SV/GCNV.  SNV_INDEL
-        # is passed as a proxy for non-SV/GCNV.
-        validate_allele_type(ht, DatasetType.SNV_INDEL)
+        validate_allele_type(ht, self.dataset_types)
         validate_no_duplicate_variants(ht, reference_genome, DatasetType.SNV_INDEL)
         # NB: we do not filter with "filter" here
         # ReferenceDatasets are DatasetType agnostic and that