From aecbd0821409e4759096d78f6e58d405f7cd7467 Mon Sep 17 00:00:00 2001
From: Kori Kuzma <korikuzma@gmail.com>
Date: Tue, 16 Jul 2024 12:43:41 -0400
Subject: [PATCH] fix: TypeError for params that accept union of enums (#566)

* Methods now raise ValueError if incorrect enum value passed
---
 src/variation/hgvs_dup_del_mode.py      | 49 +++++++++++++++++--------
 src/variation/translators/translator.py |  9 ++++-
 2 files changed, 42 insertions(+), 16 deletions(-)

diff --git a/src/variation/hgvs_dup_del_mode.py b/src/variation/hgvs_dup_del_mode.py
index fcfa8c5b..722c221e 100644
--- a/src/variation/hgvs_dup_del_mode.py
+++ b/src/variation/hgvs_dup_del_mode.py
@@ -11,6 +11,25 @@
 # Define deletion alt types
 DELS = {AltType.DELETION_AMBIGUOUS, AltType.DELETION}
 
+# Define supported alt types for HGVS Dup Del Mode
+DELS_DUPS = {
+    AltType.DELETION,
+    AltType.DELETION_AMBIGUOUS,
+    AltType.DUPLICATION,
+    AltType.DUPLICATION_AMBIGUOUS,
+}
+
+
+def _check_supported_alt_type(alt_type: AltType) -> None:
+    """Check that ``alt_type`` is one of ``DUP_DELS``
+
+    :param alt_type: Alteration type
+    :raises ValueError: If ``alt_type`` not one of ``DELS_DUPS``.
+    """
+    if alt_type not in DELS_DUPS:
+        err_msg = f"`alt_type` must be one of: {DELS_DUPS}"
+        raise ValueError(err_msg)
+
 
 class HGVSDupDelMode:
     """Class for handling how to interpret HGVS duplications and deletions."""
@@ -24,10 +43,7 @@ def __init__(self, seqrepo_access: SeqRepoAccess) -> None:
 
     def default_mode(
         self,
-        alt_type: AltType.DELETION
-        | AltType.DELETION_AMBIGUOUS
-        | AltType.DUPLICATION
-        | AltType.DUPLICATION_AMBIGUOUS,
+        alt_type: AltType,
         location: dict,
         vrs_seq_loc_ac: str,
         baseline_copies: int | None = None,
@@ -43,14 +59,17 @@ def default_mode(
         else
             allele
 
-        :param alt_type: The type of alteration
+        :param alt_type: The type of alteration. Must be one of ``DELS_DUPS``.
         :param location: Sequence Location object
         :param vrs_seq_loc_ac: Accession used in VRS Sequence Location
         :param baseline_copies: Baseline copies for Copy Number Count variation
         :param copy_change: copy change for Copy Number Change Variation
         :param alt: Alteration
+        :raises ValueError: If ``alt_type`` not one of ``DELS_DUPS``.
         :return: VRS Variation object represented as a dict
         """
+        _check_supported_alt_type(alt_type)
+
         variation = None
         if not baseline_copies and alt_type in AMBIGUOUS_REGIONS:
             variation = self.copy_number_change_mode(alt_type, location, copy_change)
@@ -62,20 +81,20 @@ def default_mode(
 
     def copy_number_count_mode(
         self,
-        alt_type: AltType.DELETION
-        | AltType.DELETION_AMBIGUOUS
-        | AltType.DUPLICATION
-        | AltType.DUPLICATION_AMBIGUOUS,
+        alt_type: AltType,
         location: dict,
         baseline_copies: int,
     ) -> dict:
         """Return a VRS Copy Number Variation.
 
-        :param alt_type: The type of alteration
+        :param alt_type: The type of alteration. Must be one of ``DELS_DUPS``.
         :param location: VRS SequenceLocation
         :param baseline_copies: Baseline copies number
+        :raises ValueError: If ``alt_type`` not one of ``DELS_DUPS``.
         :return: VRS Copy Number object represented as a dict
         """
+        _check_supported_alt_type(alt_type)
+
         copies = baseline_copies - 1 if alt_type in DELS else baseline_copies + 1
         seq_loc = models.SequenceLocation(**location)
         seq_loc.id = ga4gh_identify(seq_loc)
@@ -85,20 +104,20 @@ def copy_number_count_mode(
 
     def copy_number_change_mode(
         self,
-        alt_type: AltType.DELETION
-        | AltType.DELETION_AMBIGUOUS
-        | AltType.DUPLICATION
-        | AltType.DUPLICATION_AMBIGUOUS,
+        alt_type: AltType,
         location: dict,
         copy_change: models.CopyChange | None = None,
     ) -> dict:
         """Return copy number change variation
 
-        :param alt_type: The type of alteration
+        :param alt_type: The type of alteration. Must be one of ``DELS_DUPS``.
         :param location: VRS SequenceLocation
         :param copy_change: The copy change
+        :raises ValueError: If ``alt_type`` not one of ``DELS_DUPS``.
         :return: Copy Number Change variation as a dict
         """
+        _check_supported_alt_type(alt_type)
+
         if not copy_change:
             copy_change = (
                 models.CopyChange.EFO_0030067
diff --git a/src/variation/translators/translator.py b/src/variation/translators/translator.py
index 09ce11aa..e8ddb481 100644
--- a/src/variation/translators/translator.py
+++ b/src/variation/translators/translator.py
@@ -164,7 +164,7 @@ async def get_p_or_cdna_translation_result(
         start_pos: int,
         end_pos: int,
         alt_type: AltType,
-        coordinate_type: AnnotationLayer.PROTEIN | AnnotationLayer.CDNA,
+        coordinate_type: AnnotationLayer,
         errors: list[str],
         cds_start: int | None = None,
         ref: str | None = None,
@@ -184,8 +184,15 @@ async def get_p_or_cdna_translation_result(
             `coordinate_type == AnnotationLayer.CDNA`.
         :param ref: Expected reference sequence
         :param alt: Expected change
+        :raises ValueError: If ``coordinate`` type not one of
+            ``AnnotationLayer.PROTEIN`` or ``AnnotationLayer.CDNA``
         :return: Translation result if successful. Else, `None`
         """
+        supported_coordinate_types = {AnnotationLayer.PROTEIN, AnnotationLayer.CDNA}
+        if coordinate_type not in supported_coordinate_types:
+            err_msg = f"`coordinate_type` must be one of {supported_coordinate_types}"
+            raise ValueError(err_msg)
+
         vrs_allele = None
         vrs_seq_loc_ac = None
         vrs_seq_loc_ac_status = VrsSeqLocAcStatus.NA