tools

tools_sofa/cleanup-blastdb-volumes.py

@@ -0,0 +1,162 @@
+#!/usr/bin/env python3
+"""
+# $Id: cleanup-blastdb-volumes.py 590894 2019-08-07 14:59:53Z camacho $
+# ===========================================================================
+# 
+#                            PUBLIC DOMAIN NOTICE
+#               National Center for Biotechnology Information
+# 
+#  This software/database is a "United States Government Work" under the
+#  terms of the United States Copyright Act.  It was written as part of
+#  the author's official duties as a United States Government employee and
+#  thus cannot be copyrighted.  This software/database is freely available
+#  to the public for use. The National Library of Medicine and the U.S.
+#  Government have not placed any restriction on its use or reproduction.
+# 
+#  Although all reasonable efforts have been taken to ensure the accuracy
+#  and reliability of the software and data, the NLM and the U.S.
+#  Government do not and cannot warrant the performance or results that
+#  may be obtained by using this software or data. The NLM and the U.S.
+#  Government disclaim all warranties, express or implied, including
+#  warranties of performance, merchantability or fitness for any particular
+#  purpose.
+# 
+#  Please cite the author in any work or product based on this material.
+# 
+# ===========================================================================
+# 
+# Author:  Christiam Camacho
+# 
+# File Description:
+#   Script to remove needless BLAST database files.
+# 
+# ===========================================================================
+"""
+import argparse, os, configparser
+import unittest, tempfile
+from pathlib import Path
+from glob import glob
+
+VERSION = '1.0'
+DESC = r""" Remove needless BLAST database volumes.  """
+
+
+class Tester(unittest.TestCase):
+    """ Testing class for this script. """
+
+    def test_blastdb_config_invalid(self):
+        rv = get_blastdb_from_ncbi_config("/dev/null")
+        self.assertIsNone(rv)
+
+    def test_blastdb_config(self):
+        config = configparser.ConfigParser()
+        expected = "/blast/db/blast"
+        config['BLAST'] = {'BLASTDB': expected}
+        tf = tempfile.NamedTemporaryFile(mode="wt")
+        config.write(tf)
+        tf.flush()
+        rv = get_blastdb_from_ncbi_config(tf.name)
+        self.assertEqual(expected, rv)
+
+    def test_blastdb_finder(self):
+        tal = tempfile.NamedTemporaryFile(suffix=".pin")
+        dbname = find_blastdb(tal.name[:-4], True)
+        self.assertEqual(dbname, tal.name[:-4])
+
+
+def find_blastdb(name: str, is_prot: bool) -> str:
+    """ Returns full path to BLAST database or None. """
+    alias_file = "{}.{}al".format(name, "p" if is_prot else "n")
+    index_file = "{}.{}in".format(name, "p" if is_prot else "n")
+    if os.path.exists(alias_file) or os.path.exists(index_file):
+        return name
+
+    if "BLASTDB" in os.environ:
+        alf = os.path.join(os.environ["BLASTDB"], alias_file)
+        idxf = os.path.join(os.environ["BLASTDB"], index_file)
+        if os.path.exists(alf) or os.path.exists(idxf):
+            return os.path.join(os.environ["BLASTDB"], name)
+
+    paths = [os.getcwd(), str(Path.home())]
+    if "NCBI" in os.environ:
+        paths.append(os.path.join(os.environ["NCBI"]))
+
+    for path in paths:
+        for fname in [".ncbirc", "ncbi.ini"]:
+            ncbirc = os.path.join(path, fname)
+            if os.path.exists(ncbirc):
+                blastdb = get_blastdb_from_ncbi_config(ncbirc)
+                if blastdb is not None:
+                    alf = os.path.join(blastdb, alias_file)
+                    idxf = os.path.join(blastdb, index_file)
+                    if os.path.exists(alf) or os.path.exists(idxf):
+                        return os.path.join(blastdb, name)
+
+
+def get_blastdb_from_ncbi_config(config_file: str) -> str:
+    """ Return the BLASTDB setting from the NCBI configuration file or None. """
+    config = configparser.ConfigParser()
+    config.read(config_file)
+    if 'BLAST' in config and 'BLASTDB' in config['BLAST']:
+        return config['BLAST']['BLASTDB']
+
+
+def main():
+    """ Entry point into this program. """
+    parser = create_arg_parser()
+    args = parser.parse_args()
+
+    ext = args.dbtype[0]
+    db = find_blastdb(args.db, ext == 'p')
+    if db == None:
+        print("Cannot find {} {} BLAST database".
+              format("protein" if ext == 'p' else "nucleotide", args.db),
+              file=sys.stderr)
+        return 1
+
+    alias_file = "{}.{}al".format(db, ext)
+    if not os.path.exists(alias_file):
+        return 1
+
+    with open(alias_file, "rt") as al:
+        for line in al:
+            if not line.startswith("DBLIST"):
+                continue
+            vols = list(map(lambda x: x.replace('"', ''), line.split()[1:]))
+            for existing_vols in sorted(glob("{}.*.{}in".format(db, ext))):
+                vol_name = os.path.basename(existing_vols)[:-4]
+                if vol_name in vols:
+                    continue
+                if args.dry_run:
+                    print("Will remove extra volume {}".format(existing_vols[:-4]))
+                to_rm = glob("{}??".format(existing_vols[:-2]))
+                to_rm += glob("{}.tar.gz.md5".format(existing_vols[:-4]))
+                for f in to_rm:
+                    if not args.dry_run:
+                        os.unlink(f)
+                        print("Removed {}".format(f))
+                    elif args.verbose > 0:
+                        print("Will remove {}".format(f))
+
+    return 0
+
+
+def create_arg_parser():
+    """ Create the command line options parser object for this script. """
+    parser = argparse.ArgumentParser(description=DESC)
+    parser.add_argument("-db", required=True, help="BLAST database name")
+    parser.add_argument("-dbtype", help="Molecule type", required=True,
+                        choices=["prot", "nucl"])
+    parser.add_argument("-dry-run", action='store_true',
+                        help="Do not delete any files, just list them")
+    parser.add_argument('-version', action='version',
+                        version='%(prog)s ' + VERSION)
+    parser.add_argument("-verbose", action="count", default=0,
+                        help="Increase output verbosity")
+    return parser
+
+
+if __name__ == "__main__":
+    import sys
+
+    sys.exit(main())