163 lines
6.1 KiB
Python
163 lines
6.1 KiB
Python
#!/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())
|