[query] Add AlleleType Enum

hail/.gitignore

@@ -10,6 +10,7 @@ python/hail/docs/_build/*
 python/hail/docs/_static/hail_version.js
 python/hailtop/dist
 python/hailtop/hailctl/deploy.yaml
+python/hail/docs/genetics/hail.genetics.AlleleType.rst
 python/hail/docs/genetics/hail.genetics.Call.rst
 python/hail/docs/genetics/hail.genetics.Locus.rst
 python/hail/docs/genetics/hail.genetics.Pedigree.rst
@@ -65,6 +66,7 @@ python/hail/docs/vds/hail.vds.to_merged_sparse_mt.rst
 python/hail/docs/vds/hail.vds.local_to_global.rst
 python/hail/docs/vds/hail.vds.merge_reference_blocks.rst
 python/hail/docs/vds/hail.vds.truncate_reference_blocks.rst
+python/hail/docs/tutorials/iframe_figures/
 src/main/c/.cxx.vsn
 src/main/c/headers
 src/main/c/lib

hail/python/hail/docs/functions/genetics.rst

@@ -31,6 +31,7 @@ Genetics functions
     is_valid_locus
     contig_length
     allele_type
+    numeric_allele_type
     pl_dosage
     gp_dosage
     get_sequence
@@ -65,6 +66,7 @@ Genetics functions
 .. autofunction:: is_valid_locus
 .. autofunction:: contig_length
 .. autofunction:: allele_type
+.. autofunction:: numeric_allele_type
 .. autofunction:: pl_dosage
 .. autofunction:: gp_dosage
 .. autofunction:: get_sequence

hail/python/hail/docs/genetics/index.rst

@@ -14,6 +14,7 @@ genetics
     :toctree: ./
     :template: class.rst
 
+    AlleleType
     Call
     Locus
     Pedigree

hail/python/hail/expr/__init__.py

@@ -181,6 +181,7 @@
     is_complex,
     is_strand_ambiguous,
     allele_type,
+    numeric_allele_type,
     hamming,
     mendel_error_code,
     triangle,
@@ -406,6 +407,7 @@
     'is_complex',
     'is_strand_ambiguous',
     'allele_type',
+    'numeric_allele_type',
     'hamming',
     'mendel_error_code',
     'triangle',

hail/python/hail/expr/functions.py

@@ -85,6 +85,7 @@
     is_float32,
     is_float64,
 )
+from hail.genetics.allele_type import AlleleType
 from hail.genetics.reference_genome import reference_genome_type, ReferenceGenome
 import hail.ir as ir
 from hail.typecheck import (
@@ -3278,16 +3279,26 @@ def corr(x, y) -> Float64Expression:
     return _func("corr", tfloat64, x, y)
 
 
-_base_regex = "^([ACGTNM])+$"
-_symbolic_regex = r"(^\.)|(\.$)|(^<)|(>$)|(\[)|(\])"
-_allele_types = ["Unknown", "SNP", "MNP", "Insertion", "Deletion", "Complex", "Star", "Symbolic"]
-_allele_enum = {i: v for i, v in builtins.enumerate(_allele_types)}
-_allele_ints = {v: k for k, v in _allele_enum.items()}
-
-
 @typecheck(ref=expr_str, alt=expr_str)
 @ir.udf(tstr, tstr)
-def _num_allele_type(ref, alt) -> Int32Expression:
+def numeric_allele_type(ref, alt) -> Int32Expression:
+    """Returns the type of the polymorphism as an integer. The value returned
+    is the integer value of :class:`.AlleleType` representing that kind of
+    polymorphism.
+
+    Examples
+    --------
+
+    >>> hl.eval(hl.numeric_allele_type('A', 'T')) == AlleleType.SNP
+    True
+
+    Notes
+    -----
+    The values of :class:`.AlleleType` are not stable and thus should not be
+    relied upon across hail versions.
+    """
+    _base_regex = "^([ACGTNM])+$"
+    _symbolic_regex = r"(^\.)|(\.$)|(^<)|(>$)|(\[)|(\])"
     return hl.bind(
         lambda r, a: hl.if_else(
             r.matches(_base_regex),
@@ -3299,24 +3310,33 @@ def _num_allele_type(ref, alt) -> Int32Expression:
                     r.length() == a.length(),
                     hl.if_else(
                         r.length() == 1,
-                        hl.if_else(r != a, _allele_ints['SNP'], _allele_ints['Unknown']),
-                        hl.if_else(hamming(r, a) == 1, _allele_ints['SNP'], _allele_ints['MNP']),
+                        hl.if_else(r != a, AlleleType.SNP, AlleleType.UNKNOWN),
+                        hl.if_else(hamming(r, a) == 1, AlleleType.SNP, AlleleType.MNP),
                     ),
                 )
-                .when((r.length() < a.length()) & (r[0] == a[0]) & a.endswith(r[1:]), _allele_ints["Insertion"])
-                .when((r[0] == a[0]) & r.endswith(a[1:]), _allele_ints["Deletion"])
-                .default(_allele_ints['Complex']),
+                .when((r.length() < a.length()) & (r[0] == a[0]) & a.endswith(r[1:]), AlleleType.INSERTION)
+                .when((r[0] == a[0]) & r.endswith(a[1:]), AlleleType.DELETION)
+                .default(AlleleType.COMPLEX),
             )
-            .when(a == '*', _allele_ints['Star'])
-            .when(a.matches(_symbolic_regex), _allele_ints['Symbolic'])
-            .default(_allele_ints['Unknown']),
-            _allele_ints['Unknown'],
+            .when(a == '*', AlleleType.STAR)
+            .when(a.matches(_symbolic_regex), AlleleType.SYMBOLIC)
+            .default(AlleleType.UNKNOWN),
+            AlleleType.UNKNOWN,
         ),
         ref,
         alt,
     )
 
 
+@deprecated(version='0.2.129', reason="Replaced by the public numeric_allele_type")
+@typecheck(ref=expr_str, alt=expr_str)
+def _num_allele_type(ref, alt) -> Int32Expression:
+    """Provided for backwards compatibility, don't use it in new code, or
+    within the hail library itself
+    """
+    return numeric_allele_type(ref, alt)
+
+
 @typecheck(ref=expr_str, alt=expr_str)
 def is_snp(ref, alt) -> BooleanExpression:
     """Returns ``True`` if the alleles constitute a single nucleotide polymorphism.
@@ -3338,7 +3358,7 @@ def is_snp(ref, alt) -> BooleanExpression:
     -------
     :class:`.BooleanExpression`
     """
-    return _num_allele_type(ref, alt) == _allele_ints["SNP"]
+    return numeric_allele_type(ref, alt) == AlleleType.SNP
 
 
 @typecheck(ref=expr_str, alt=expr_str)
@@ -3362,7 +3382,7 @@ def is_mnp(ref, alt) -> BooleanExpression:
     -------
     :class:`.BooleanExpression`
     """
-    return _num_allele_type(ref, alt) == _allele_ints["MNP"]
+    return numeric_allele_type(ref, alt) == AlleleType.MNP
 
 
 @typecheck(ref=expr_str, alt=expr_str)
@@ -3458,7 +3478,7 @@ def is_insertion(ref, alt) -> BooleanExpression:
     -------
     :class:`.BooleanExpression`
     """
-    return _num_allele_type(ref, alt) == _allele_ints["Insertion"]
+    return numeric_allele_type(ref, alt) == AlleleType.INSERTION
 
 
 @typecheck(ref=expr_str, alt=expr_str)
@@ -3482,7 +3502,7 @@ def is_deletion(ref, alt) -> BooleanExpression:
     -------
     :class:`.BooleanExpression`
     """
-    return _num_allele_type(ref, alt) == _allele_ints["Deletion"]
+    return numeric_allele_type(ref, alt) == AlleleType.DELETION
 
 
 @typecheck(ref=expr_str, alt=expr_str)
@@ -3506,9 +3526,7 @@ def is_indel(ref, alt) -> BooleanExpression:
     -------
     :class:`.BooleanExpression`
     """
-    return hl.bind(
-        lambda t: (t == _allele_ints["Insertion"]) | (t == _allele_ints["Deletion"]), _num_allele_type(ref, alt)
-    )
+    return hl.bind(lambda t: (t == AlleleType.INSERTION) | (t == AlleleType.DELETION), numeric_allele_type(ref, alt))
 
 
 @typecheck(ref=expr_str, alt=expr_str)
@@ -3532,7 +3550,7 @@ def is_star(ref, alt) -> BooleanExpression:
     -------
     :class:`.BooleanExpression`
     """
-    return _num_allele_type(ref, alt) == _allele_ints["Star"]
+    return numeric_allele_type(ref, alt) == AlleleType.STAR
 
 
 @typecheck(ref=expr_str, alt=expr_str)
@@ -3556,7 +3574,7 @@ def is_complex(ref, alt) -> BooleanExpression:
     -------
     :class:`.BooleanExpression`
     """
-    return _num_allele_type(ref, alt) == _allele_ints["Complex"]
+    return numeric_allele_type(ref, alt) == AlleleType.COMPLEX
 
 
 @typecheck(ref=expr_str, alt=expr_str)
@@ -3624,7 +3642,7 @@ def allele_type(ref, alt) -> StringExpression:
     -------
     :class:`.StringExpression`
     """
-    return hl.literal(_allele_types)[_num_allele_type(ref, alt)]
+    return hl.literal(AlleleType.strings())[numeric_allele_type(ref, alt)]
 
 
 @typecheck(s1=expr_str, s2=expr_str)

hail/python/hail/genetics/__init__.py

@@ -1,6 +1,7 @@
+from .allele_type import AlleleType
 from .call import Call
 from .reference_genome import ReferenceGenome
 from .pedigree import Pedigree, Trio
 from .locus import Locus
 
-__all__ = ['Locus', 'Call', 'Pedigree', 'Trio', 'ReferenceGenome']
+__all__ = ['AlleleType', 'Locus', 'Call', 'Pedigree', 'Trio', 'ReferenceGenome']

hail/python/hail/genetics/allele_type.py

@@ -0,0 +1,96 @@
+from enum import IntEnum, auto
+
+
+_ALLELE_STRS = (
+    "Unknown",
+    "SNP",
+    "MNP",
+    "Insertion",
+    "Deletion",
+    "Complex",
+    "Star",
+    "Symbolic",
+    "Transition",
+    "Transversion",
+)
+
+
+class AlleleType(IntEnum):
+    """An enumeration for allele type.
+
+    Notes
+    -----
+    The precise values of the enumeration constants are not guarenteed
+    to be stable and must not be relied upon.
+    """
+
+    UNKNOWN = 0
+    """Unknown Allele Type"""
+    SNP = auto()
+    """Single-nucleotide Polymorphism (SNP)"""
+    MNP = auto()
+    """Multi-nucleotide Polymorphism (MNP)"""
+    INSERTION = auto()
+    """Insertion"""
+    DELETION = auto()
+    """Deletion"""
+    COMPLEX = auto()
+    """Complex Polymorphism"""
+    STAR = auto()
+    """Star Allele (``alt=*``)"""
+    SYMBOLIC = auto()
+    """Symbolic Allele
+
+    e.g. ``alt=<INS>``
+    """
+    TRANSITION = auto()
+    """Transition SNP
+
+    e.g. ``ref=A alt=G``
+
+    Note
+    ----
+    This is only really used internally in :func:`hail.vds.sample_qc` and
+    :func:`hail.methods.sample_qc`.
+    """
+    TRANSVERSION = auto()
+    """Transversion SNP
+
+    e.g. ``ref=A alt=C``
+
+    Note
+    ----
+    This is only really used internally in :func:`hail.vds.sample_qc` and
+    :func:`hail.methods.sample_qc`.
+    """
+
+    def __str__(self):
+        return str(self.value)
+
+    @property
+    def pretty_name(self):
+        """A formatted (as opposed to uppercase) version of the member's name,
+        to match :func:`~hail.expr.functions.allele_type`
+
+        Examples
+        --------
+        >>> AlleleType.INSERTION.pretty_name
+        'Insertion'
+        >>> at = AlleleType(hl.eval(hl.numeric_allele_type('a', 'att')))
+        >>> at.pretty_name == hl.eval(hl.allele_type('a', 'att'))
+        True
+        """
+        return _ALLELE_STRS[self]
+
+    @classmethod
+    def _missing_(cls, value):
+        if not isinstance(value, str):
+            return None
+        return cls.__members__.get(value.upper())
+
+    @staticmethod
+    def strings():
+        """Returns the names of the allele types, for use with
+        :func:`~hail.expr.functions.literal`
+        """
+        return list(_ALLELE_STRS)