bowtie2-build

#!/usr/bin/env python3

"""
 Copyright 2014, Ben Langmead <langmea@cs.jhu.edu>

 This file is part of Bowtie 2.

 Bowtie 2 is free software: you can redistribute it and/or modify
 it under the terms of the GNU General Public License as published by
 the Free Software Foundation, either version 3 of the License, or
 (at your option) any later version.

 Bowtie 2 is distributed in the hope that it will be useful,
 but WITHOUT ANY WARRANTY; without even the implied warranty of
 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 GNU General Public License for more details.

 You should have received a copy of the GNU General Public License
 along with Bowtie 2.  If not, see <http://www.gnu.org/licenses/>.
"""

import os
import sys
import gzip
import inspect
import logging
import argparse
import subprocess


def get_gunzip_size(fn):
    size = 0
    with gzip.open(fn) as f:
        while True:
            data = f.read(8192)
            size += len(data)
            if not data:
                break
    return size

def sais_enabled(build_bin_spec):
    output = subprocess.check_output([build_bin_spec, "--version"], universal_newlines=True)

    return output.find("USE_SAIS") != -1

def main():
    parser = argparse.ArgumentParser(add_help = False)

    parser.add_argument('--large-index', action='store_true')
    parser.add_argument('--verbose', action='store_true')

    group = parser.add_mutually_exclusive_group()
    group.add_argument('--debug', action='store_true')
    group.add_argument('--sanitized', action='store_true')

    logging.basicConfig(level=logging.ERROR,
                        format='%(levelname)s: %(message)s'
                        )
    delta               = 200
    small_index_max_size= 4 * 1024**3 - delta
    build_bin_name      = 'bowtie2-build'
    build_bin_s         = 'bowtie2-build-s'
    build_bin_l         = 'bowtie2-build-l'
    curr_script         = os.path.realpath(inspect.getsourcefile(main))
    ex_path             = os.path.dirname(curr_script)
    build_bin_spec      = os.path.join(ex_path,build_bin_s)

    if sais_enabled(build_bin_spec):
        small_index_max_size = 2 * 1024 ** 3 - delta

    script_options, argv = parser.parse_known_args()

    if script_options.verbose:
        logging.getLogger().setLevel(logging.INFO)

    if script_options.debug:
        build_bin_spec += "-debug"
        build_bin_l += "-debug"

    if script_options.sanitized:
        build_bin_spec += "-sanitized"
        build_bin_l += "-sanitized"

    fastas = []
    if "-c" not in argv and len(argv) >= 2:
        for index in reversed(range(len(argv) - 1)):
            arg = argv[index]
            if arg.startswith("-") or arg.isdigit():
                break
            fastas.insert(0, arg)
            argv.remove(arg)
        if fastas:
            fastas = ",".join(fastas)
            argv.insert(len(argv) - 1, fastas)

    if script_options.large_index:
        build_bin_spec = os.path.join(ex_path, build_bin_l)
    elif fastas:
        tot_size = 0
        for fn in fastas.split(","):
            if os.path.exists(fn):
                if fn.endswith(".gz"):
                    tot_size += get_gunzip_size(fn)
                elif fn.endswith(".zst"):
                    sys.stderr.write(
                        "{} cannot determine the uncompressed size of ZSTD compressed files.\
                        Please make sure that the uncompressed file size is suitable for\
                        building the current index".format(
                            curr_script
                        )
                    )
                else:
                    statinfo = os.stat(fn)
                    tot_size += statinfo.st_size
        if tot_size > small_index_max_size:
            build_bin_spec = os.path.join(ex_path, build_bin_l)

    if not os.path.exists(build_bin_spec):
        sys.stderr.write(
            "{0} does not exist, try running `[g]make {0}'\n".format(
                os.path.basename(build_bin_spec)
            )
        )
        sys.exit(1)

    argv.insert(0, "basic-0")
    argv.insert(0, "--wrapper")
    argv.insert(0, build_bin_spec)
    logging.info("Command: %s" % " ".join(argv))
    sys.exit(subprocess.call(argv))


if __name__ == "__main__":
    main()