Merge pull request #210 from pbashyal-nmdp/cli_improvements

Cli Tools improvements

Author: mmaiers-nmdp

pyard/data_repository.py

@@ -486,7 +486,9 @@ def generate_short_nulls(db_connection, who_group):
     return shortnulls
 
 
-def generate_mac_codes(db_connection: sqlite3.Connection, refresh_mac: bool):
+def generate_mac_codes(
+    db_connection: sqlite3.Connection, refresh_mac: bool = False, load_mac: bool = True
+):
     """
     MAC files come in 2 different versions:
 
@@ -530,29 +532,31 @@ def generate_mac_codes(db_connection: sqlite3.Connection, refresh_mac: bool):
 
     :param db_connection: Database connection to the sqlite database
     :param refresh_mac: Refresh the database with newer MAC data ?
+    :param load_mac: Should MAC be loaded at all
     :return: None
     """
-    mac_table_name = "mac_codes"
-    if refresh_mac or not db.table_exists(db_connection, mac_table_name):
-        # Load the MAC file to a DataFrame
-        mac_url = "https://hml.nmdp.org/mac/files/numer.v3.zip"
-        df_mac = pd.read_csv(
-            mac_url,
-            sep="\t",
-            compression="zip",
-            skiprows=3,
-            names=["Code", "Alleles"],
-            keep_default_na=False,
-        )
-        # Create a dict from code to alleles
-        mac = df_mac.set_index("Code")["Alleles"].to_dict()
-        # Save the mac dict to db
-        db.save_dict(
-            db_connection,
-            table_name=mac_table_name,
-            dictionary=mac,
-            columns=("code", "alleles"),
-        )
+    if load_mac:
+        mac_table_name = "mac_codes"
+        if refresh_mac or not db.table_exists(db_connection, mac_table_name):
+            # Load the MAC file to a DataFrame
+            mac_url = "https://hml.nmdp.org/mac/files/numer.v3.zip"
+            df_mac = pd.read_csv(
+                mac_url,
+                sep="\t",
+                compression="zip",
+                skiprows=3,
+                names=["Code", "Alleles"],
+                keep_default_na=False,
+            )
+            # Create a dict from code to alleles
+            mac = df_mac.set_index("Code")["Alleles"].to_dict()
+            # Save the mac dict to db
+            db.save_dict(
+                db_connection,
+                table_name=mac_table_name,
+                dictionary=mac,
+                columns=("code", "alleles"),
+            )
 
 
 def to_serological_name(locus_name: str):

pyard/db.py

@@ -55,7 +55,7 @@ def create_db_connection(data_dir, imgt_version, ro=False):
         # Open the database in read-only mode
         file_uri = f"file:{db_filename}?mode=ro"
         # Multiple threads can access the same connection since it's only ro
-        return sqlite3.connect(file_uri, check_same_thread=False, uri=True)
+        return sqlite3.connect(file_uri, check_same_thread=False, uri=True), db_filename
 
     # Check the imgt_version is a valid IMGT DB Version
     # by querying the IMGT site
@@ -73,7 +73,7 @@ def create_db_connection(data_dir, imgt_version, ro=False):
 
     # Open the database for read/write
     file_uri = f"file:{db_filename}"
-    return sqlite3.connect(file_uri, uri=True)
+    return sqlite3.connect(file_uri, uri=True), db_filename
 
 
 def table_exists(connection: sqlite3.Connection, table_name: str) -> bool:

pyard/misc.py

@@ -94,3 +94,14 @@ def get_data_dir(data_dir):
     else:
         data_dir = db.get_pyard_db_install_directory()
     return data_dir
+
+
+def get_imgt_version(imgt_version):
+    if imgt_version:
+        version = imgt_version.replace(".", "")
+        if version.isdigit():
+            return version
+        raise RuntimeError(
+            f"{imgt_version} is not a valid IMGT database version number"
+        )
+    return "Latest"

pyard/pyard.py

@@ -80,7 +80,11 @@ class ARD(object):
     """
 
     def __init__(
-        self, imgt_version: str = "Latest", data_dir: str = None, config: dict = None
+        self,
+        imgt_version: str = "Latest",
+        data_dir: str = None,
+        load_mac: bool = True,
+        config: dict = None,
     ):
         """
         ARD will load valid alleles, xx codes and MAC mappings for the given
@@ -97,10 +101,10 @@ def __init__(
             self._config.update(config)
 
         # Create a database connection for writing
-        self.db_connection = db.create_db_connection(data_dir, imgt_version)
+        self.db_connection, _ = db.create_db_connection(data_dir, imgt_version)
 
         # Load MAC codes
-        dr.generate_mac_codes(self.db_connection, False)
+        dr.generate_mac_codes(self.db_connection, refresh_mac=False, load_mac=load_mac)
         # Load ARS mappings
         self.ars_mappings, p_group = dr.generate_ars_mapping(
             self.db_connection, imgt_version
@@ -140,7 +144,7 @@ def __init__(
             gc.freeze()
 
         # Re-open the connection in read-only mode as we're not updating it anymore
-        self.db_connection = db.create_db_connection(data_dir, imgt_version, ro=True)
+        self.db_connection, _ = db.create_db_connection(data_dir, imgt_version, ro=True)
 
     def __del__(self):
         """
@@ -393,7 +397,7 @@ def validate(self, glstring):
         except InvalidAlleleError as e:
             raise InvalidTypingError(
                 f"{glstring} is not valid GL String. \n {e.message}", e
-            )
+            ) from None
 
     def is_XX(self, glstring: str, loc_antigen: str = None, code: str = None) -> bool:
         if loc_antigen is None or code is None:
@@ -718,7 +722,7 @@ def refresh_mac_codes(self) -> None:
         Refreshes MAC code for the current IMGT db version.
         :return: None
         """
-        dr.generate_mac_codes(self.db_connection, True)
+        dr.generate_mac_codes(self.db_connection, refresh_mac=True)
 
     def get_db_version(self) -> str:
         """

scripts/pyard

@@ -25,18 +25,8 @@ import argparse
 import sys
 
 import pyard
-
-
-def get_imgt_version(imgt_version):
-    if imgt_version:
-        version = imgt_version.replace(".", "")
-        if version.isdigit():
-            return version
-        raise RuntimeError(
-            f"{imgt_version} is not a valid IMGT database version number"
-        )
-    return None
-
+from pyard.exceptions import InvalidAlleleError
+from pyard.misc import get_data_dir, get_imgt_version
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(
@@ -51,6 +41,12 @@ if __name__ == "__main__":
         action="store_true",
         help="IPD-IMGT/HLA DB Version number",
     )
+    parser.add_argument(
+        "-d",
+        "--data-dir",
+        dest="data_dir",
+        help="Data directory to store imported data",
+    )
     parser.add_argument(
         "-i",
         "--imgt-version",
@@ -68,22 +64,33 @@ if __name__ == "__main__":
 
     args = parser.parse_args()
 
+    if args.splits:
+        mapping = pyard.find_broad_splits(args.splits)
+        if mapping:
+            print(f"{mapping[0]} = {'/'.join(mapping[1])}")
+        sys.exit(0)
+
     imgt_version = get_imgt_version(args.imgt_version)
-    if imgt_version:
-        ard = pyard.ARD(imgt_version)
-    else:
-        ard = pyard.ARD()
+    data_dir = get_data_dir(args.data_dir)
+    ard = pyard.ARD(imgt_version=imgt_version, data_dir=data_dir)
 
     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)
 
-    if args.splits:
-        mapping = pyard.find_broad_splits(args.splits)
-        if mapping:
-            print(f"{mapping[0]} = {'/'.join(mapping[1])}")
-        sys.exit(0)
+    try:
+        if args.redux_type:
+            print(ard.redux_gl(args.gl_string, args.redux_type))
+        else:
+            for redux_type in pyard.pyard.reduction_types:
+                redux_type_info = f"Reduction Method: {redux_type}"
+                print(redux_type_info)
+                print("-" * len(redux_type_info))
+                print(ard.redux_gl(args.gl_string, redux_type))
+    except InvalidAlleleError as e:
+        print("Error:", e)
 
-    print(ard.redux_gl(args.gl_string, args.redux_type))
+    # Remove ard and close db connection
     del ard

scripts/pyard-import

@@ -68,11 +68,13 @@ if __name__ == "__main__":
         help="Show Versions of available IMGT Databases",
     )
     parser.add_argument(
-        "--db-version",
+        "-i",
+        "--imgt-version",
         dest="imgt_version",
         help="Import supplied IMGT_VERSION DB Version",
     )
     parser.add_argument(
+        "-d",
         "--data-dir",
         dest="data_dir",
         help="Data directory to store imported data",
@@ -92,6 +94,12 @@ if __name__ == "__main__":
         action="store_true",
         help="reinstall a fresh version of database",
     )
+    parser.add_argument(
+        "--skip-mac",
+        dest="skip_mac",
+        action="store_true",
+        help="Skip creating MAC mapping",
+    )
     args = parser.parse_args()
 
     if args.show_versions:
@@ -118,8 +126,14 @@ if __name__ == "__main__":
             db_fullname.unlink(missing_ok=True)
 
     print(f"Importing IMGT database version: {imgt_version}")
+    if args.skip_mac:
+        load_mac = False
+        print(f"Skipping MAC tables creation")
+    else:
+        load_mac = True
+
     try:
-        ard = pyard.ARD(imgt_version=imgt_version, data_dir=data_dir)
+        ard = pyard.ARD(imgt_version=imgt_version, data_dir=data_dir, load_mac=load_mac)
     except ValueError as e:
         print(f"Error importing version {imgt_version}:", e)
         sys.exit(1)
@@ -141,6 +155,6 @@ if __name__ == "__main__":
 
     if args.refresh_mac:
         print(f"Updating MACs")
-        db_connection = db.create_db_connection(data_dir, imgt_version, ro=False)
-        data_repository.generate_mac_codes(db_connection, True)
+        db_connection, _ = db.create_db_connection(data_dir, imgt_version, ro=False)
+        data_repository.generate_mac_codes(db_connection, refresh_mac=True)
         print(f"Updated MACs for {imgt_version} IMGT database.")

scripts/pyard-reduce-csv

@@ -38,6 +38,7 @@ import pandas as pd
 import pyard
 import pyard.drbx as drbx
 from pyard.exceptions import PyArdError
+from pyard.misc import get_data_dir, get_imgt_version
 
 
 def is_serology(allele: str) -> bool:
@@ -171,11 +172,21 @@ def create_drbx(row, locus_in_allele_name):
 
 
 if __name__ == "__main__":
-
     # config is specified with a -c parameter
     parser = argparse.ArgumentParser()
     parser.add_argument("-c", "--config", help="JSON Configuration file", required=True)
-
+    parser.add_argument(
+        "-d",
+        "--data-dir",
+        dest="data_dir",
+        help="Data directory to store imported data",
+    )
+    parser.add_argument(
+        "-i",
+        "--imgt-version",
+        dest="imgt_version",
+        help="IPD-IMGT/HLA db to use for redux",
+    )
     args = parser.parse_args()
     config_filename = args.config
 
@@ -197,8 +208,9 @@ if __name__ == "__main__":
             print("  pip install openpyxl")
             sys.exit(1)
 
-    # Instantiate py-ard object with the latest
-    ard = pyard.ARD()
+    data_dir = get_data_dir(args.data_dir)
+    imgt_version = get_imgt_version(args.imgt_version)
+    ard = pyard.ARD(imgt_version=imgt_version, data_dir=data_dir)
 
     # Read the Input File
     # Read only the columns to be saved.

scripts/pyard-status

@@ -39,21 +39,25 @@ def get_latest_imgt_version() -> int:
     return max(map(int, pyard.db_versions()[:-1]))
 
 
+def get_file_size(file_name: str) -> float:
+    return os.path.getsize(file_name) / 1024 / 1024
+
+
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(
         description="""
         py-ard tool to provide a status report for reference SQLite databases.
-        """,
+        """
     )
     parser.add_argument(
+        "-d",
         "--data-dir",
         dest="data_dir",
         help="Data directory to store imported data",
     )
 
     args = parser.parse_args()
     data_dir = get_data_dir(args.data_dir)
-    # print(data_dir)
 
     imgt_regex = re.compile(r"pyard-(.+)\.sqlite3")
     for _, _, filenames in os.walk(data_dir):
@@ -62,7 +66,9 @@ if __name__ == "__main__":
             # eg: get 3440 from 'pyard-3440.sqlite3'
             match = imgt_regex.match(filename)
             imgt_version = match.group(1)  # Get first group
-            db_connection = db.create_db_connection(data_dir, imgt_version, ro=True)
+            db_connection, db_filename = db.create_db_connection(
+                data_dir, imgt_version, ro=True
+            )
             print("-" * 43)
             if imgt_version == "Latest":
                 db_version = data_repository.get_db_version(db_connection)
@@ -80,11 +86,16 @@ if __name__ == "__main__":
                     )
             else:
                 print(f"IMGT DB Version: {imgt_version}")
+            file_size = get_file_size(db_filename)
+            print(f"File: {db_filename}")
+            print(f"Size: {file_size:.2f}MB")
             print("-" * 43)
             print(f"|{'Table Name':20}|{'Rows':20}|")
             print(f"|{'-' * 41}|")
             for table in (
-                data_repository.ars_mapping_tables + data_repository.code_mapping_tables
+                data_repository.ars_mapping_tables
+                + data_repository.code_mapping_tables
+                + ["mac_codes"]
             ):
                 if db.table_exists(db_connection, table):
                     total_rows = db.count_rows(db_connection, table)