- Allowpyard and pyard-reduce-csv to support --data-dir and --imgt-version options

- pyard-status shows the sqlite db and the corresponding file size

Author: pbashyal-nmdp

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

@@ -101,7 +101,7 @@ 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, refresh_mac=False, load_mac=load_mac)
@@ -144,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):
         """

scripts/pyard

@@ -26,18 +26,7 @@ import sys
 
 import pyard
 from pyard.exceptions import InvalidAlleleError
-
-
-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.misc import get_data_dir, get_imgt_version
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(
@@ -52,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",
@@ -69,24 +64,22 @@ 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))

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",
@@ -153,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)
+        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,6 +86,9 @@ 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}|")