Add type hints to assemblymapper.py

Author: anju24

src/hgvs/assemblymapper.py

@@ -5,6 +5,7 @@
 from bioutils.sequences import TranslationTable
 
 import hgvs
+from hgvs.alignmentmapper import AlignmentMapper
 import hgvs.normalizer
 from hgvs.exceptions import (
     HGVSDataNotAvailableError,
@@ -13,6 +14,8 @@
     HGVSUnsupportedOperationError,
 )
 from hgvs.variantmapper import VariantMapper
+from hgvs.sequencevariant import SequenceVariant
+from typing import Optional
 
 _logger = logging.getLogger(__name__)
 
@@ -50,7 +53,7 @@ class AssemblyMapper(VariantMapper):
 
     def __init__(
         self,
-        hdp,
+        hdp: hgvs.dataproviders.interface.Interface,
         assembly_name=hgvs.global_config.mapping.assembly,
         alt_aln_method=hgvs.global_config.mapping.alt_aln_method,
         normalize=hgvs.global_config.mapping.normalize,
@@ -60,7 +63,7 @@ def __init__(
         add_gene_symbol=hgvs.global_config.mapping.add_gene_symbol,
         *args,
         **kwargs,
-    ):
+    ) -> None:
         """
         :param object hdp: instance of hgvs.dataprovider subclass
         :param bool replace_reference: replace reference (entails additional network access)
@@ -88,44 +91,54 @@ def __init__(
         self.in_par_assume = in_par_assume
         self._norm = None
         if self.normalize:
-            vm = VariantMapper(hdp=hdp, replace_reference=replace_reference,
-                               prevalidation_level=prevalidation_level,
-                               add_gene_symbol=add_gene_symbol)
+            vm = VariantMapper(
+                hdp=hdp,
+                replace_reference=replace_reference,
+                prevalidation_level=prevalidation_level,
+                add_gene_symbol=add_gene_symbol,
+            )
             self._norm = hgvs.normalizer.Normalizer(
-                hdp, alt_aln_method=alt_aln_method, validate=False, variantmapper=vm,
+                hdp,
+                alt_aln_method=alt_aln_method,
+                validate=False,
+                variantmapper=vm,
             )
         self._assembly_map = {
-            k: v for k, v in hdp.get_assembly_map(self.assembly_name).items() if k.startswith("NC_")
+            k: v
+            for k, v in hdp.get_assembly_map(self.assembly_name).items()
+            if k.startswith("NC_")
         }
         self._assembly_accessions = set(self._assembly_map.keys())
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return (
             "{self.__module__}.{t.__name__}(alt_aln_method={self.alt_aln_method}, "
             "assembly_name={self.assembly_name}, normalize={self.normalize}, "
             "prevalidation_level={self.prevalidation_level}, "
-            "replace_reference={self.replace_reference})".format(self=self, t=type(self))
+            "replace_reference={self.replace_reference})".format(
+                self=self, t=type(self)
+            )
         )
 
-    def g_to_c(self, var_g, tx_ac):
+    def g_to_c(self, var_g: SequenceVariant, tx_ac: str) -> SequenceVariant:
         var_out = super(AssemblyMapper, self).g_to_c(
             var_g, tx_ac, alt_aln_method=self.alt_aln_method
         )
         return self._maybe_normalize(var_out)
 
-    def g_to_n(self, var_g, tx_ac):
+    def g_to_n(self, var_g: SequenceVariant, tx_ac: str) -> SequenceVariant:
         var_out = super(AssemblyMapper, self).g_to_n(
             var_g, tx_ac, alt_aln_method=self.alt_aln_method
         )
         return self._maybe_normalize(var_out)
 
-    def g_to_t(self, var_g, tx_ac):
+    def g_to_t(self, var_g: SequenceVariant, tx_ac: str) -> SequenceVariant:
         var_out = super(AssemblyMapper, self).g_to_t(
             var_g, tx_ac, alt_aln_method=self.alt_aln_method
         )
         return self._maybe_normalize(var_out)
 
-    def c_to_g(self, var_c):
+    def c_to_g(self, var_c: SequenceVariant) -> SequenceVariant:
         alt_ac = self._alt_ac_for_tx_ac(var_c.ac)
         var_out = super(AssemblyMapper, self).c_to_g(
             var_c, alt_ac, alt_aln_method=self.alt_aln_method
@@ -146,7 +159,7 @@ def t_to_g(self, var_t):
         )
         return self._maybe_normalize(var_out)
 
-    def t_to_p(self, var_t):
+    def t_to_p(self, var_t: SequenceVariant) -> SequenceVariant:
         """Return a protein variant, or "non-coding" for non-coding variant types
 
         CAUTION: Unlike other x_to_y methods that always return
@@ -165,36 +178,44 @@ def t_to_p(self, var_t):
             "Expected a coding (c.) or non-coding (n.) variant; got " + str(var_t)
         )
 
-    def c_to_n(self, var_c):
+    def c_to_n(self, var_c: SequenceVariant) -> SequenceVariant:
         alt_ac = self._alt_ac_for_tx_ac(var_c.ac)
         var_out = super(AssemblyMapper, self).c_to_n(
             var_c, alt_ac=alt_ac, alt_aln_method=self.alt_aln_method
         )
         return self._maybe_normalize(var_out)
 
-    def n_to_c(self, var_n):
+    def n_to_c(self, var_n: SequenceVariant) -> SequenceVariant:
         alt_ac = self._alt_ac_for_tx_ac(var_n.ac)
         var_out = super(AssemblyMapper, self).n_to_c(
             var_n, alt_ac=alt_ac, alt_aln_method=self.alt_aln_method
         )
         return self._maybe_normalize(var_out)
 
-    def c_to_p(self, var_c, translation_table=TranslationTable.standard):
+    def c_to_p(
+        self, var_c: SequenceVariant, translation_table=TranslationTable.standard
+    ) -> SequenceVariant:
         alt_ac = self._alt_ac_for_tx_ac(var_c.ac)
         var_out = super(AssemblyMapper, self).c_to_p(
-            var_c, alt_ac=alt_ac, alt_aln_method=self.alt_aln_method, translation_table=translation_table
+            var_c,
+            alt_ac=alt_ac,
+            alt_aln_method=self.alt_aln_method,
+            translation_table=translation_table,
         )
         return self._maybe_normalize(var_out)
 
-    def relevant_transcripts(self, var_g):
+    def relevant_transcripts(self, var_g: SequenceVariant) -> SequenceVariant:
         """return list of transcripts accessions (strings) for given variant,
         selected by genomic overlap"""
         tx = self.hdp.get_tx_for_region(
-            var_g.ac, self.alt_aln_method, var_g.posedit.pos.start.base, var_g.posedit.pos.end.base
+            var_g.ac,
+            self.alt_aln_method,
+            var_g.posedit.pos.start.base,
+            var_g.posedit.pos.end.base,
         )
         return [e["tx_ac"] for e in tx]
 
-    def _alt_ac_for_tx_ac(self, tx_ac):
+    def _alt_ac_for_tx_ac(self, tx_ac: str) -> str:
         """return chromosomal accession for given transcript accession (and
         the_assembly and aln_method setting used to instantiate this
         AssemblyMapper)
@@ -219,12 +240,18 @@ def _alt_ac_for_tx_ac(self, tx_ac):
             names = set(self._assembly_map[ac] for ac in alt_acs)
             if names != set("XY"):
                 alts = ", ".join(
-                    ["{ac} ({n})".format(ac=ac, n=self._assembly_map[ac]) for ac in alt_acs]
+                    [
+                        "{ac} ({n})".format(ac=ac, n=self._assembly_map[ac])
+                        for ac in alt_acs
+                    ]
                 )
                 raise HGVSError(
                     "Multiple chromosomal alignments for {tx_ac} in {an}"
                     " using {am} (non-pseudoautosomal region) [{alts}]".format(
-                        tx_ac=tx_ac, an=self.assembly_name, am=self.alt_aln_method, alts=alts
+                        tx_ac=tx_ac,
+                        an=self.assembly_name,
+                        am=self.alt_aln_method,
+                        alts=alts,
                     )
                 )
 
@@ -237,7 +264,9 @@ def _alt_ac_for_tx_ac(self, tx_ac):
                     )
                 )
 
-            alt_acs = [ac for ac in alt_acs if self._assembly_map[ac] == self.in_par_assume]
+            alt_acs = [
+                ac for ac in alt_acs if self._assembly_map[ac] == self.in_par_assume
+            ]
             if len(alt_acs) != 1:
                 raise HGVSError(
                     "Multiple chromosomal alignments for {tx_ac} in {an}"
@@ -253,7 +282,12 @@ def _alt_ac_for_tx_ac(self, tx_ac):
         assert len(alt_acs) == 1, "Should have exactly one alignment at this point"
         return alt_acs[0]
 
-    def _fetch_AlignmentMapper(self, tx_ac, alt_ac=None, alt_aln_method=None):
+    def _fetch_AlignmentMapper(
+        self,
+        tx_ac: str,
+        alt_ac: Optional[str] = None,
+        alt_aln_method: Optional[str] = None,
+    ) -> AlignmentMapper:
         """convenience version of VariantMapper._fetch_AlignmentMapper that
         derives alt_ac from transcript, assembly, and alt_aln_method
         used to instantiate the AssemblyMapper instance
@@ -264,9 +298,11 @@ def _fetch_AlignmentMapper(self, tx_ac, alt_ac=None, alt_aln_method=None):
             alt_ac = self._alt_ac_for_tx_ac(tx_ac)
         if alt_aln_method is None:
             alt_aln_method = self.alt_aln_method
-        return super(AssemblyMapper, self)._fetch_AlignmentMapper(tx_ac, alt_ac, alt_aln_method)
+        return super(AssemblyMapper, self)._fetch_AlignmentMapper(
+            tx_ac, alt_ac, alt_aln_method
+        )
 
-    def _maybe_normalize(self, var):
+    def _maybe_normalize(self, var: SequenceVariant) -> SequenceVariant:
         """normalize variant if requested, and ignore HGVSUnsupportedOperationError
         This is better than checking whether the variant is intronic because
         future UTAs will support LRG, which will enable checking intronic variants.