Merge pull request #281 from pbashyal-nmdp/feature/similar

FInd similar alleles/MACs

Author: mmaiers-nmdp

pyard/db.py

@@ -356,16 +356,16 @@ def load_dict(
     return table_as_dict
 
 
-def similar_alleles(connection: sqlite3.Connection, allele_name: str) -> Set[str]:
+def similar_alleles(connection: sqlite3.Connection, allele_prefix: str) -> Set[str]:
     """
-    Find similar alleles starting with the provided allele_name.
+    Find similar alleles starting with the provided prefix.
 
     :param connection: db connection of type sqlite.Connection
-    :param allele_name: Allele name to use as a prefix to find similar alleles
+    :param allele_prefix: Allele name to use as a prefix to find similar alleles
     :return: list of similar alleles
     """
     query = "SELECT allele FROM alleles WHERE allele LIKE ?"
-    cursor = connection.execute(query, (f"{allele_name}%",))
+    cursor = connection.execute(query, (f"{allele_prefix}%",))
     result = cursor.fetchall()
     # fetchall() returns a list of tuples of results
     # e.g. [('C*04:09N',)]
@@ -374,6 +374,24 @@ def similar_alleles(connection: sqlite3.Connection, allele_name: str) -> Set[str
     return alleles
 
 
+def similar_mac(connection: sqlite3.Connection, mac_prefix: str) -> Set[str]:
+    """
+    Find similar MAC codes starting with the provided prefix.
+
+    :param connection: db connection of type sqlite.Connection
+    :param mac_prefix: MAC fragment to use as a prefix to find similar MACs
+    :return: list of similar MAC codes
+    """
+    query = "SELECT code FROM mac_codes WHERE code LIKE ?"
+    cursor = connection.execute(query, (f"{mac_prefix}%",))
+    result = cursor.fetchall()
+    # fetchall() returns a list of tuples of results
+    # e.g. [('DJZUP',)]
+    # Get out the first value of the tuple from the result list
+    codes = set(map(lambda t: t[0], result))
+    return codes
+
+
 def find_serology_for_allele(
     connection: sqlite3.Connection, allele_name: str
 ) -> Dict[str, str]:

scripts/pyard

@@ -22,13 +22,92 @@
 #    > http://www.opensource.org/licenses/lgpl-license.php
 #
 import argparse
+import functools
 import sys
 
+from pyard import smart_sort
 from pyard.constants import VALID_REDUCTION_TYPES
 import pyard.misc
+from pyard.db import similar_alleles, similar_mac
 from pyard.exceptions import InvalidAlleleError, InvalidTypingError, InvalidMACError
 from pyard.misc import get_data_dir, get_imgt_version
 
+
+def find_similar_alleles(ard, prefix):
+    if "*" in prefix:  # Only for those that have locus
+        locus, fields = prefix.split("*")
+        if fields:  # Only if at least a field is specified after *
+            if len(fields.split(":")) == 2:  # Check for MACs
+                first_field, mac_prefix = fields.split(":")
+                if mac_prefix.isalpha():
+                    similar_mac_names = similar_mac(ard.db_connection, mac_prefix)
+                    if similar_mac_names:
+                        locus_prefix = f"{locus}*{first_field}"
+                        # TODO: validate all the mac codes with the prefix
+                        #       show only the valid macs
+                        for code in sorted(similar_mac_names):
+                            print(f"{locus_prefix}:{code}")
+        else:
+            # Nothing after *
+            sys.exit(2)
+    else:
+        # No *
+        sys.exit(1)
+
+    # find similar alleles
+    similar_allele_names = similar_alleles(ard.db_connection, prefix)
+    if similar_allele_names:
+        for allele in sorted(
+            similar_allele_names,
+            key=functools.cmp_to_key(smart_sort.smart_sort_comparator),
+        ):
+            print(allele)
+    sys.exit(0)
+
+
+def lookup_mac_codes():
+    global e
+    try:
+        mac = ard.lookup_mac(args.lookup_mac)
+        print(mac)
+    except InvalidMACError as e:
+        print(e.message, file=sys.stderr)
+    sys.exit(0)
+
+
+def expand_mac_code():
+    global allele_list, e
+    try:
+        allele_list = ard.expand_mac(args.expand_mac)
+        print(allele_list)
+    except InvalidMACError as e:
+        print(e.message, file=sys.stderr)
+    sys.exit(0)
+
+
+def find_broad_splits():
+    mapping = pyard.find_broad_splits(args.splits)
+    if mapping:
+        print(f"{mapping[0]} = {'/'.join(mapping[1])}")
+    sys.exit(0)
+
+
+def show_version():
+    version = ard.get_db_version()
+    print(f"IPD-IMGT/HLA version:", version)
+    print(f"py-ard version:", pyard.__version__)
+    sys.exit(0)
+
+
+def perform_cwd_redux():
+    global cwd_redux
+    if args.validate:
+        ard.validate(args.cwd)
+    cwd_redux = ard.cwd_redux(args.cwd)
+    print(cwd_redux)
+    sys.exit(0)
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(
         description="""
@@ -76,6 +155,11 @@ if __name__ == "__main__":
     parser.add_argument(
         "--lookup-mac", dest="lookup_mac", help="Lookup MAC for an Allele List"
     )
+    parser.add_argument(
+        "--similar",
+        dest="similar_allele",
+        help="Find Similar Alleles with given prefix",
+    )
     parser.add_argument(
         "--non-strict",
         dest="non_strict",
@@ -100,45 +184,32 @@ if __name__ == "__main__":
 
     ard = pyard.init(imgt_version=imgt_version, data_dir=data_dir, config=new_config)
 
+    # Handle --version option
     if args.version:
-        version = ard.get_db_version()
-        print(f"IPD-IMGT/HLA version:", version)
-        print(f"py-ard version:", pyard.__version__)
-        sys.exit(0)
+        show_version()
 
+    # Handle --splits option
     if args.splits:
-        mapping = pyard.find_broad_splits(args.splits)
-        if mapping:
-            print(f"{mapping[0]} = {'/'.join(mapping[1])}")
-        sys.exit(0)
+        find_broad_splits()
 
     # Handle --expand-mac option
     if args.expand_mac:
-        try:
-            allele_list = ard.expand_mac(args.expand_mac)
-            print(allele_list)
-        except InvalidMACError as e:
-            print(e.message, file=sys.stderr)
-        sys.exit(0)
+        expand_mac_code()
 
     # Handle --lookup-mac option
     if args.lookup_mac:
-        try:
-            mac = ard.lookup_mac(args.lookup_mac)
-            print(mac)
-        except InvalidMACError as e:
-            print(e.message, file=sys.stderr)
-        sys.exit(0)
+        lookup_mac_codes()
+
+    # Handle --similar option
+    if args.similar_allele:
+        find_similar_alleles(ard, args.similar_allele)
 
     try:
+        if args.cwd:
+            perform_cwd_redux()
+
         if args.validate and args.gl_string:
             ard.validate(args.gl_string)
-        if args.cwd:
-            if args.validate:
-                ard.validate(args.cwd)
-            cwd_redux = ard.cwd_redux(args.cwd)
-            print(cwd_redux)
-            sys.exit(0)
 
         if args.redux_type:
             print(ard.redux(args.gl_string, args.redux_type))