def md5(fname, string=False):
    """straight from quantumsoup
    http://stackoverflow.com/questions/3431825/
        generating-an-md5-checksum-of-a-file
    updated 20160916 for strings; if true, fname can be a string
    """
    if string:
        return hashlib.md5(fname.encode('utf-8')).hexdigest()
    else:
        hash_md5 = hashlib.md5()
        with open(fname, "rb") as f:
            for chunk in iter(lambda: f.read(4096), b""):
                hash_md5.update(chunk)
            return hash_md5.hexdigest() def parse_fasta_header(first_line):
    """ return accession from first line of a fasta file as a string

    Args:
        first_line (str): The first line of a file
    Raise:
        ValueError: if first_line doesnt start with a ">"
        ValueError: if regex to extract accession fails to get a group
    Returns:
        str: name of the accession

    """
    if not first_line.startswith(">"):
        raise ValueError("A valid fasta must start with a '>'!")
    if first_line.startswith(">gi"):
        accession = re.search(r">gi\|.*?\|.*?\|(.*?)\|",
                              first_line).groups()[0]
    else:
        try:  # this handles normal flavored fastas with '>seqid descr'
            accession = re.search(r">(.*?) .*", first_line).groups()[0]
        except AttributeError:  # this handles headers with only '>seqid\n'
            accession = first_line.strip().replace(">", "")
    if accession is None:
        raise ValueError(
            str("unable to extract accession from first line in file: \n" +
                first_line))
    return accession def add_locus_tags_to_gff(gff, acc):
    """ Add fake locus tags to gff file in place.

    Args:
        gff (str): path to gff file
        acc (str): accession to label the loci with
    Returns:
        retrun path to new gff file (the old one, but with _tagged at the end

    """
    LOCUS = 0
    gff_list = []
    with open(gff, 'r') as g:
        for idx, line in enumerate(g):
            line = line.strip()
            if idx == 0:
                gff_list.append(line)
            else:
                gff_list.append(
                    re.sub("Name=",
                           "locus_tag={0}_{1};Name=".format(acc, LOCUS),
                           line))
                LOCUS = LOCUS + 1

    new_gff = str(os.path.splitext(gff)[0] + "_tagged.gff")
    with open(new_gff, 'w') as outgff:
        for l in gff_list:
            # print(str(l))
            outgff.write(str(l + "\n"))
    return(new_gff) def combine_gbs(finalgb, gb_list):
    """ pretty self explanatory - it combines files

    it skips lines starting with "##"

    Args:
        finalgb (str): path to ourput file
        gb_list (str): list of other file paths to be combined
    Returns:
        None: doesn't return something

    """
    with open(finalgb, 'w') as outfile:
        for idx, fname in enumerate(gb_list):
            with open(fname) as infile:
                for line in infile:
                    if line.startswith("##") and idx != 0:
                        continue
                    outfile.write(line) def append_accession_and_version(accession, ingb, finalgb):
    """ add accession, version and GI to a gb file

    Biopython requires a version and an accession to parse correctly.
    This adds it

    Args:
        accession (str): Accession string to add
        ingb (str): path to initial genbank file
        finalgb (str): path to resulting genbank file
    Returns:
        None
    Raises:
        None

    """
    with open(finalgb, 'w') as outfile:
        with open(ingb) as infile:
            for idx, line in enumerate(infile):
                if idx == 1:
                    repline = "ACCESSION   {0}\nVERSION     {0}{1}{2}".format(
                        accession,
                        "  GI:0000000\n",
                        line)
                    outfile.write(repline)
                else:
                    outfile.write(line) def splitMultifasta(multi, output, name, dirname="contigs", logger=None):
    """create new files containing a single fasta entry each

    regex stolen from SO
    name is the name of the file, output is the parent dir for the output dir

    Args:
        multi (str): path to multifasta
        output (str): path to output directory
        name (str): accession for renaming multifasta
        dirname (str): name of directory to fold split fastas
    Returns:
        None
    Raises:
        None

    """
    idlist = []
    assert logger is not None, "must use logging"
    os.makedirs(os.path.join(output, dirname))
    #source_fname = os.path.splitext(os.path.basename(multi))[0]
    with open(os.path.expanduser(multi), "r") as mf:
        for idx, rec in enumerate(SeqIO.parse(mf, "fasta")):
            logger.debug("record %d: %s", idx, rec.id)
            # logger.debug(rec)
            if name is None:
                name = rec.id
            # if fname in idlist:
            fname = name + "_" + str(idx)
            rec.id  = fname
            if not re.match("^[a-zA-Z0-9_\.]*$", fname):
                logger.warning("Problem with header %s", fname)
                logger.warning("file header contains special characters! " +
                               "Valid characters are in set [a-zA-Z0-9_]; " +
                               "rename with the --name option to prevent " +
                               "later issues. Renaming as source filename " +
                               "with index, sorr..." )
            if len(rec.id) > 16:
                logger.warning("Fixing problem with header %s: must be shorter than 16 chars", fname)

            with open(os.path.join(output, dirname,
                                   fname + ".fa"), "w") as outf:
                renamed_rec = SeqRecord(rec.seq, id=rec.id,
                                        description="from riboScan")
                SeqIO.write(renamed_rec, outf, "fasta")
            idlist.append(fname)
    logger.debug("rewrote the following records")
    logger.debug("\n".join(idlist)) def getFastas(inp, output_root, name, logger):
    """return list of fasta files

    If the input is a multifasta, a new directory is created to contain
    single sequence files.  given an input that could either
    be a single fasta, multifasta,
    or a directory, retrun the appropriate files

    Args:return list of fasta files
        inp (str): either the path to a file or a directory
        output_root: path to where the contigs dir will be created, if needed
        name (str): accession for resulting contig files or sequences
    Returns:
        (list): paths to sequences
    Raises:
        SystemExit: if file is empty or path doesnt exist
        SystemExit: if no files end in *fasta or *fa

    """
    if not os.path.isdir(os.path.expanduser(inp)):
        if not os.path.isfile(os.path.expanduser(inp)):
            logger.error("'%s' is not a valid directory or file!",
                         os.path.expanduser(inp))
            sys.exit(1)
        if not os.path.getsize(os.path.expanduser(inp)) > 0:
            logger.error("'%s' looks like an empty file!",
                         os.path.expanduser(inp))
            sys.exit(1)
        logger.info("spliting multifasta into multiple fastas " +
                    "for easier processing")
        splitMultifasta(multi=inp, output=output_root, name=name,
                        logger=logger)

        fastas = glob.glob(os.path.join(output_root, "contigs",
                                        "*.fa"))
    else:
        fastas = glob.glob(os.path.join(os.path.expanduser(inp),
                                        "*.fa"))
    if len(fastas) == 0:
        logger.info("No fasta files found with extention '.fa'. " +
                       "searching for '*.fasta' files instead...")
        fastas = glob.glob(os.path.join(os.path.expanduser(inp),
                                        "*.fasta"))
        if len(fastas) == 0:  # still
            logger.error("No files in %s with extention '*fasta'! Exiting",
                         inp)
            sys.exit(1)
    # unify the output of multifasta
    combine_gbs(finalgb=os.path.join(output_root, "scannedScaffolds.fa"),
                gb_list=fastas)
    return(fastas) def checkSingleFasta(fasta, logger):
    """raise systen exit if multiple entries in fasta.

    This is mostly depreciated, but left in place juuuust in case

    Args:
        fasta (type): path to file
    Returns:
        None
    Raises:
        SystemExit: thrown if fasta has multiple entries

    """
    with open(fasta, 'r') as f:
        counter = 0
        for rec in SeqIO.parse(f, "fasta"):
            counter = counter + 1
            if counter > 1:
                logger.error(
                    "Fasta has multiple entries! " +
                    "use splitMultifasta.sh to create single-entry " +
                    "fastas from it. exiting...")
                sys.exit(1) def main(args, logger=None):
    # allow user to give relative paths
    output_root = os.path.abspath(os.path.expanduser(args.output))
    sys_barrnap = shutil.which(args.barrnap_exe)
    assert sys_barrnap is not None, \
        "barrnap executable not found!"
    test_barrnap_ok(exe=sys_barrnap)
    try:
        os.makedirs(output_root, exist_ok=False)
    except OSError:
        print("Output directory already exists; exiting...")
        sys.exit(1)
    t0 = time.time()
    if logger is None:
        logger = set_up_logging(
            verbosity=args.verbosity,
            outfile=os.path.join(output_root, "riboScan.log"),
            name=__name__)
    logger.info("Usage:\n%s\n", " ".join([x for x in sys.argv]))
    logger.debug("All settings used:")
    for k, v in sorted(vars(args).items()):
        logger.debug("%s: %s", k, str(v))

    output_file = os.path.join(output_root, "scannedScaffolds.gb")
    output_file_gff = os.path.join(output_root, "scannedScaffolds.gff")
    ##  get and check list of input files
    fastas = getFastas(inp=args.contigs, output_root=output_root,
                       name=args.name, logger=logger)
    gb_list = []
    gff_list = []  # for tagged gffs, that is
    for fasta in sorted(fastas):
        # check for multiple entry fastas.
        # This can still happen if mutlifastas are in dir input
        checkSingleFasta(fasta, logger=logger)
        with open(fasta, 'r') as f:
            rec = next(SeqIO.parse(f, "fasta"))
        accession = rec.id
        if args.min_length > len(rec.seq):
            logger.debug("Skipping %s: must be %d bp, set by --min_length",
                         accession, args.min_length)
            continue
        logger.info("Accession for %s: %s", fasta, accession)
        barrnap_cmd = make_barrnap_cmd(
            infasta=fasta,
            outgff=os.path.join(output_root, "{0}.gff".format(accession)),
            exe=sys_barrnap,
            threads=args.cores,
            thresh=args.id_thresh,
            kingdom=args.kingdom)
        logger.info("running barrnap cmd: %s", barrnap_cmd)
        subprocess.run(barrnap_cmd,
                       shell=sys.platform != "win32",
                       stdout=subprocess.PIPE,
                       stderr=subprocess.PIPE,
                       check=True)
        logger.info("grooming gff")
        tagged_gff = add_locus_tags_to_gff(
            gff=os.path.join(output_root, "{0}.gff".format(accession)),
            acc=accession)
        unfinished_gb = os.path.join(
            output_root, "{0}_pre.gb".format(accession))
        logger.info("creating GenBank file")
        try:
            make_genbank(
                fasta=fasta,
                gff=tagged_gff,
                outpath=os.path.join(output_root, "{0}_pre.gb".format(accession)))
        except Exception as e:
            logger.error(e)
            sys.exit(1)
        append_accession_and_version(
            accession=accession,
            ingb=unfinished_gb,
            finalgb=os.path.join(output_root, "{0}.gb".format(accession)))
        gb_list.append(os.path.join(output_root, "{0}.gb".format(accession)))
        gff_list.append(tagged_gff)
    ### no more fastas
    combine_gbs(finalgb=output_file, gb_list=gb_list)
    # this is why you give things intellegent method names
    combine_gbs(finalgb=output_file_gff, gb_list=gff_list)
    # Report that we've finished
    logger.info("Done: %s", time.asctime())
    logger.info("combined scaffolds can be found here: %s", output_file)
    logger.info("Time taken: %.3fs" % (time.time() - t0)) def make_barrnap_cmd(infasta, outgff, exe, thresh, kingdom, evalue=1e-06, threads=1):
    """ construct a system call to barrnap

    Args:
        infasta (str): path to fasta file
        outgff (str): name of gff file to create
        exe (str): path to system executable for barrnap
        thresh (float): Proportional length threshold to reject prediction
        kingdom (str): bac, euk, arc, or mito
        threads (int): number of threads to use
    Returns:
        (str): command

    """
    assert 1 >= thresh > 0 and thresh , "Thresh must be between 0 and 1!"
    if exe.endswith("py"):
        # ensure running python barrnap uses >3.5
        pyexe = str(sys.executable + " ")
    else:
        pyexe = ""
    cmd = "{0}{1} -k {2} {3} --reject {4} --threads {5} --evalue {6} > {7}".format(
        pyexe,
        exe,
        kingdom,
        infasta,
        thresh,
        threads,
        evalue,
        outgff
        )
    return cmd
import sys
import logging
import shutil
import time
import os
import unittest
from Bio.SeqRecord import SeqRecord
from argparse import Namespace
from unittest.mock import MagicMock, patch
from riboSeed.shared_methods import md5
from riboSeed.riboScan import parse_fasta_header, \
    add_locus_tags_to_gff, combine_gbs, append_accession_and_version, \
    splitMultifasta, getFastas, checkSingleFasta, main
from riboSeed.shared_methods import make_barrnap_cmd
# -*- coding: utf-8 -*-
"""
Created on Tue Aug 30 08:57:31 2016
@author: nicholas

"""
# import subprocess
# import multiprocessing

# from Bio.Seq import Seq


sys.dont_write_bytecode = True

logger = logging


@unittest.skipIf((sys.version_info[0] != 3) or (sys.version_info[1] < 5),
                 "Subprocess.call among other things wont run if tried " +
                 " with less than python 3.5")
class riboSeedTestCase(unittest.TestCase):
    """ tests for riboSeed.py
    """
    def setUp(self):
        self.test_dir = os.path.join(os.path.dirname(__file__),
                                     "output_riboScan_tests")
        self.ref_dir = os.path.join(os.path.dirname(__file__), "references")
        self.scan_ref_dir = os.path.join(os.path.dirname(__file__),
                                         "references",
                                         "riboScan_references")
        self.no_locus_gff = os.path.join(self.scan_ref_dir,
                                         "no_locus.gff")
        self.with_locus_gff = os.path.join(self.scan_ref_dir,
                                           "with_locus.gff")
        self.combined_file = os.path.join(self.scan_ref_dir,
                                          "combined_gbs_just_kidding.txt")
        self.multifasta = os.path.join(self.scan_ref_dir,
                                       "multiFasta.fasta")
        self.no_accession_gb = os.path.join(self.scan_ref_dir,
                                            'no_accession_or_version.gb')
        self.with_accession_gb = os.path.join(self.scan_ref_dir,
                                              'with_accession_or_version.gb')

        self.good_contig = os.path.join(self.ref_dir,
                                        'contigs.fasta')
        self.short_contig = os.path.join(self.ref_dir,
                                         'contigs.fasta')
        self.args = Namespace(skip_contol=False, kmers="21,33,55,77,99",
                              spades_exe="spades.py",
                              quast_exe="python2.7 quast.py",
                              cores=2)
        self.startTime = time.time()
        self.cores = 2
        self.maxDiff = 2000
        self.to_be_removed = []

    def test_gi_parse_fasta_header(self):
        """check valid fasta header parsing"""
        header1 = str(">gi|218703261|ref|NC_011751.1| " +
                      "Escherichia coli UMN026 chromosome, complete genome\n")
        self.assertEqual("NC_011751.1", parse_fasta_header(header1))

    def test_nongi_parse_fasta_header(self):
        """check nonstandard fasta header parsing"""
        header2 = str(">testgenome gi|218703261|ref|NC_011751.1| " +
                      "Escherichia coli UMN026 chromosome, complete genome\n")
        self.assertEqual("testgenome", parse_fasta_header(header2))

    def test_fail_parse_fasta_header(self):
        """ensudre bad fastas fail header parsing"""
        header3 = str("testgenome|blabla|")
        with self.assertRaises(ValueError):
            parse_fasta_header(header3)

    def test_make_barrnap_cmd(self):
        """ check barrnap command construction  """
        cmd1 = make_barrnap_cmd(infasta="test.fasta", outgff="test.gff",
                                exe="barrnap.exe", thresh=.2, kingdom='euk')
        ref_cmd1 = str("barrnap.exe -k euk test.fasta --reject 0.2 --threads" +
                       " 1 --evalue 1e-06 > test.gff")
        self.assertEqual(cmd1, ref_cmd1)

    def test_fail_barrnap_exe_main(self):
        """ fail main with bad barrnap exe"""
        with self.assertRaises(AssertionError):
            main(Namespace(contigs="test.fasta",
                           output="test",
                           barrnap_exe="definitelynotbarrnap",
                           cores=2))

    # @unittest.skipIf(shutil.which("barrnap") is None,
    #                  "barrnap executable not found. If this isnt an " +
    #                  "error from travis deployment, you probably " +
    #                  "should install it")
    def test_fail_thresh_make_barrnap_cmd(self):
        with self.assertRaises(AssertionError):
            make_barrnap_cmd(infasta="test.fasta", outgff="test.gff",
                             exe="barrnap", thresh=1.2,
                             kingdom='euk')

    def test_add_locus_tags_to_gff(self):
        """check create proper gff """
        add_locus_tags_to_gff(gff=self.no_locus_gff,
                              acc="BA000007.2")
        new_gff = str(os.path.splitext(self.no_locus_gff)[0] + "_tagged.gff")

        self.assertEqual(md5(new_gff), md5(self.with_locus_gff))