Standalone Python re-implementation of the POLCA polisher from MaSuRCA 
#!/usr/bin/env python3
"""pypolca"""

from pathlib import Path

import click
from loguru import logger

from pypolca.utils.fix_consensus_from_vcf import fix_consensus_from_vcf
from pypolca.utils.mapping import (
    bam_to_sorted_bam,
    bwa_index,
    bwa_paired,
    bwa_single,
    sam_to_bam,
    samtools_index,
)
from pypolca.utils.report import create_report
from pypolca.utils.util import (
    begin_pypolca,
    copy_file,
    end_pypolca,
    get_version,
    print_citation,
    remove_directory,
)
from pypolca.utils.validation import (
    check_dependencies,
    check_memory_limit,
    instantiate_dirs,
    validate_fasta,
    validate_fastq,
)
from pypolca.utils.variants import run_freebayes, samtools_faidx

"""
some code adapted from tbpore https://github.com/mbhall88/tbpore
"""

log_fmt = (
    "[<green>{time:YYYY-MM-DD HH:mm:ss}</green>] <level>{level: <8}</level> | "
    "<level>{message}</level>"
)


def common_options(func):
    """Common command line args
    Define common command line args here, and include them with the @common_options decorator below.
    """
    options = [
        click.option(
            "-a",
            "--assembly",
            help="Path to assembly contigs or scaffolds.",
            type=click.Path(),
            required=True,
        ),
        click.option(
            "-1",
            "--reads1",
            required=True,
            type=click.Path(),
            help="Path to polishing reads R1 FASTQ. Can be FASTQ or FASTQ gzipped. Required.",
        ),
        click.option(
            "-2",
            "--reads2",
            default=None,
            type=click.Path(),
            help="Path to polishing reads R2 FASTQ. Can be FASTQ or FASTQ gzipped. Optional. Only use -1 if you have single end reads.",
        ),
        click.option(
            "-t",
            "--threads",
            help="Number of threads.",
            default=1,
            show_default=True,
        ),
        click.option(
            "-o",
            "--output",
            default="output_polca",
            show_default=True,
            type=click.Path(),
            help="Output directory path",
        ),
        click.option(
            "-f", "--force", is_flag=True, help="Force overwrites the output directory"
        ),
        click.option(
            "-n",
            "--no_polish",
            is_flag=True,
            help="do not polish, just create vcf file, evaluate the assembly and exit",
        ),
        click.option(
            "-m",
            "--memory_limit",
            type=str,
            default="2G",
            help="Memory per thread to use in samtools sort, set to 2G or more for large genomes",
            show_default=True,
        ),
        click.option(
            "-p",
            "--prefix",
            type=str,
            default="polca",
            help="prefix for output files",
            show_default=True,
        ),
    ]
    for option in reversed(options):
        func = option(func)
    return func


@click.group()
@click.help_option("--help", "-h")
@click.version_option(get_version(), "--version", "-V")
def main_cli():
    1 + 1


"""
run command
"""


@main_cli.command()
@click.help_option("--help", "-h")
@click.version_option(get_version(), "--version", "-V")
@click.pass_context
@common_options
def run(
    ctx,
    assembly,
    reads1,
    reads2,
    threads,
    output,
    force,
    no_polish,
    memory_limit,
    prefix,
    **kwargs,
):
    """Python implementation of the POLCA polisher from MaSuRCA"""

    params = {
        "--assembly": assembly,
        "--reads1": reads1,
        "--reads2": reads2,
        "--output": output,
        "--threads": threads,
        "--force": force,
        "--memory_limit": memory_limit,
        "--no_polish": no_polish,
        "--prefix": prefix,
    }

    # validates the directory  (need to before I start or else no log file is written)
    instantiate_dirs(output, force)

    output: Path = Path(output)

    temp_dir: Path = Path(output) / "temp"
    temp_dir.mkdir(parents=True, exist_ok=True)
    logdir: Path = Path(output) / "logs"
    logdir.mkdir(parents=True, exist_ok=True)

    # begin pypolca - initial logging etc
    start_time = begin_pypolca(params)

    # check dependencies
    check_dependencies()

    # validates fasta
    validate_fasta(assembly)

    # validating FASTQ
    # whether paired or single
    # paired_flag is True
    if reads2 is None:
        logger.warning("You have not specified -2 or --reads2.")
        logger.warning("pypolca will proceed using single end reads only.")
        paired_flag = False
        validate_fastq(reads1)
    else:
        paired_flag = True
        validate_fastq(reads1)
        validate_fastq(reads2)

    logger.info(f"Checking memory limit of {memory_limit}.")
    check_memory_limit(memory_limit)

    #########
    #
    logger.info("Creating BWA index")
    assembly_temp: Path = Path(temp_dir) / "assembly.fasta"

    copy_file(Path(assembly), assembly_temp)
    bwa_index(assembly_temp, logdir)

    logger.info("Aligning reads with BWA")
    sam: Path = temp_dir / "temp_bwa.sam"
    if paired_flag is True:
        bwa_paired(reads1, reads2, assembly_temp, sam, threads, logdir)
    else:
        bwa_single(reads1, assembly_temp, sam, threads, logdir)

    logger.info("Sorting and indexing alignment file")
    bam: Path = temp_dir / "temp_bwa.bam"
    sorted_bam: Path = temp_dir / "temp_bwa_sorted.bam"
    sam_to_bam(sam, bam, threads, logdir)
    bam_to_sorted_bam(bam, sorted_bam, threads, memory_limit, logdir)
    samtools_index(sorted_bam, logdir)

    logger.info("Calling variants.")
    samtools_faidx(assembly_temp, logdir)
    vcf: Path = Path(output) / f"{prefix}.vcf"
    run_freebayes(assembly_temp, sorted_bam, output, vcf, logdir)

    ####
    # fix_consensus_from_vcf

    out_fasta: Path = Path(output) / f"{prefix}_corrected.fasta"
    if no_polish is False:
        fix_consensus_from_vcf(assembly_temp, vcf, out_fasta)

    ######
    # report.py
    report_file: Path = Path(output) / f"{prefix}.report"
    create_report(vcf, assembly_temp, report_file)

    # finalise
    if no_polish is False:
        logger.info(
            f"Final report is in {report_file}; polished assembly is in {out_fasta}"
        )
    else:
        logger.info(f"Final report is in {report_file}")

    # cleanup temp
    remove_directory(temp_dir)

    # end pypolca
    end_pypolca(start_time)


@click.command()
def citation(**kwargs):
    """Print the citations for polca"""
    print_citation()


main_cli.add_command(citation)


def main():
    main_cli()


if __name__ == "__main__":
    main()

src/pypolca/__init__.py
gbouras13/pypolca
import hashlib
import shlex
import subprocess
import sys
from pathlib import Path
from typing import List, Optional, Tuple

import click
from loguru import logger


class ExternalTool:
    def __init__(
        self,
        tool: str,
        input: str,
        output: str,
        params: str,
        logdir: Path,
        outfile: Path,
    ):
        self.command: List[str] = self._build_command(tool, input, output, params)
        logdir.mkdir(parents=True, exist_ok=True)
        command_hash = hashlib.sha256(self.command_as_str.encode("utf-8")).hexdigest()
        tool_name = Path(tool).name
        logfile_prefix: Path = logdir / f"{tool_name}_{command_hash}"
        self.out_log = f"{logfile_prefix}.out"
        self.err_log = f"{logfile_prefix}.err"
        self.outfile = outfile
        self.tool_str = tool

    @property
    def command_as_str(self) -> str:
        # return shlex.join(self.command)
        # shlex.join only for python 3.8 and on
        # freebayes 1.3.1 needs python 3.7
        return " ".join(self.command)

    @staticmethod
    def _build_command(tool: str, input: str, output: str, params: str) -> List[str]:
        # note: shlex.join does not allow us to shlex.split() later
        # this is explicitly a " ".join()
        command = " ".join([tool, params, output, input])
        escaped_command = shlex.split(command)
        return escaped_command

    def run(self) -> None:
        with open(self.out_log, "w") as stdout_fh, open(self.err_log, "w") as stderr_fh:
            print(f"Command line: {self.command_as_str}", file=stderr_fh)
            logger.info(f"Started running {self.command_as_str} ...")
            self._run_core(self.command, stdout_fh=stdout_fh, stderr_fh=stderr_fh)
            logger.info(f"Done running {self.command_as_str}")

    def run_to_stdout(
        self,
    ) -> None:
        with open(self.outfile, "w") as outfile, open(self.err_log, "w") as stderr_fh:
            print(f"Command line: {self.command_as_str}", file=stderr_fh)
            logger.info(f"Started running {self.command_as_str} ...")
            self._run_core(self.command, stdout_fh=outfile, stderr_fh=stderr_fh)
            logger.info(f"Done running {self.command_as_str}")

    @staticmethod
    def _run_core(command: List[str], stdout_fh, stderr_fh) -> None:
        subprocess.check_call(command, stdout=stdout_fh, stderr=stderr_fh)

    @staticmethod
    def run_tools(
        tools_to_run: Tuple["ExternalTool", ...], ctx: Optional[click.Context] = None
    ) -> None:
        for tool in tools_to_run:
            try:
                tool.run()
            except subprocess.CalledProcessError as error:
                logger.error(
                    f"Error calling {tool.command_as_str} (return code {error.returncode})"
                )
                logger.error(f"Please check stdout log file: {tool.out_log}")
                logger.error(f"Please check stderr log file: {tool.err_log}")
                logger.error("Temporary files are preserved for debugging")
                logger.error("Exiting...")

                if ctx:
                    ctx.exit(1)
                else:
                    sys.exit(1)

    """
    Only one tool
    """

    @staticmethod
    def run_tool(
        tool: "ExternalTool",
        ctx: Optional[click.Context] = None,
        to_stdout: Optional[bool] = False,
    ) -> None:
        try:
            if to_stdout is False:
                tool.run()
            else:  # if tool needs to write to stdout
                tool.run_to_stdout()
        except subprocess.CalledProcessError as error:
            logger.warning(
                f"Error calling {tool.command_as_str} (return code {error.returncode})"
            )
            logger.warning(f"Please check stdout log file: {tool.out_log}")
            logger.warning(f"Please check stderr log file: {tool.err_log}")
            logger.warning("Temporary files are preserved for debugging")
            logger.error("Exiting...")

            if ctx:
                ctx.exit(1)
            else:
                sys.exit(1)

src/pypolca/utils/external_tools.py
gbouras13/pypolca
#!/usr/bin/env python3

from pathlib import Path
from typing import Dict

from Bio import SeqIO
from Bio.Seq import Seq
from loguru import logger

from pypolca.utils.util import copy_file


def fix_consensus_from_vcf(ref_contigs: Path, vcf: Path, out_fasta: Path) -> None:
    """
    Fix errors in the consensus called in a VCF file by FreeBayes.

    Args:
        ref_contigs (Path): Path to the reference contigs in FASTA format.
        vcf (Path): Path to the VCF file containing variants.
        out_fasta (Path): Path to the output FASTA file.

    Returns:
        None
    """

    rseq = {}

    # Read reference sequences into memory
    with open(ref_contigs) as fasta_file:
        for record in SeqIO.parse(fasta_file, "fasta"):
            rseq[record.id] = str(record.seq)

    # Initialize variables
    ctg = ""
    fixes = []
    originals = []
    offsets = []
    fixed_sequences = {}

    # Read and process VCF file
    with open(vcf) as vcf_file:
        for line in vcf_file:
            line = line.strip()
            if line.startswith("#"):
                continue
            fields = line.split()
            if "," in fields[4] or fields[0] not in rseq:
                continue
            if fields[0] != ctg:
                if fixes:
                    if ctg not in rseq:
                        raise Exception(
                            "Sequence {} not found in the input fasta file".format(ctg)
                        )
                    oldseq = rseq[ctg]

                    # Proceed with fixing
                    for i in range(
                        len(fixes) - 1, -1, -1
                    ):  # Going in reverse order to avoid shifting sequence due to indels
                        # First, we check if the sequence at the given offset matches the original variant
                        original_seq = oldseq[
                            offsets[i] - 1 : offsets[i] - 1 + len(originals[i])
                        ]
                        if (
                            set(original_seq).intersection("acgtnACGTN")
                            and not originals[i].upper() == original_seq.upper()
                        ):
                            logger.warning(
                                "WARNING! Sequence does not match the original {} {} {} {}".format(
                                    ctg, original_seq, originals[i], offsets[i]
                                )
                            )
                        else:
                            # Then, substitute
                            oldseq = (
                                oldseq[: offsets[i] - 1]
                                + fixes[i]
                                + oldseq[offsets[i] - 1 + len(originals[i]) :]
                            )

                    fixed_sequences[ctg] = oldseq

                fixes = []
                originals = []
                offsets = []
                ctg = fields[0]

            ff = fields[9].split(":")
            if int(ff[5]) > 1 and int(ff[5]) >= 2 * int(ff[3]):
                fixes.append(fields[4])
                originals.append(fields[3])
                offsets.append(int(fields[1]))

    if fixes:
        logger.info(f"POLCA has found variants. Fixing")
        # Proceed with fixing for the last contig
        if ctg not in rseq:
            raise Exception("Sequence {} not found in the input fasta file".format(ctg))
        oldseq = rseq[ctg]

        for i in range(
            len(fixes) - 1, -1, -1
        ):  # Going in reverse order to avoid shifting sequence due to indels
            # First, we check if the sequence at the given offset matches the original variant
            original_seq = oldseq[offsets[i] - 1 : offsets[i] - 1 + len(originals[i])]
            if (
                set(original_seq).intersection("acgtnACGTN")
                and not originals[i].upper() == original_seq.upper()
            ):
                logger.warning(
                    "WARNING! Sequence does not match the original {} {} {} {}".format(
                        ctg, original_seq, originals[i], offsets[i]
                    )
                )
            else:
                # Then, substitute
                oldseq = (
                    oldseq[: offsets[i] - 1]
                    + fixes[i]
                    + oldseq[offsets[i] - 1 + len(originals[i]) :]
                )

        fixed_sequences[ctg] = oldseq

        records = []
        for contig, sequence in fixed_sequences.items():
            record = SeqIO.SeqRecord(seq=Seq(sequence), id=contig, description="")
            records.append(record)

        with open(out_fasta, "w") as output_file:
            SeqIO.write(records, output_file, "fasta")

    else:
        logger.warning(f"POLCA has found 0 variants.")
        logger.warning(
            f"The corrected FASTA {out_fasta} will be the same as the input FASTA {ref_contigs}."
        )
        # copy the reference to the output
        copy_file(ref_contigs, out_fasta)

src/pypolca/utils/fix_consensus_from_vcf.py
gbouras13/pypolca
from pathlib import Path

from pypolca.utils.external_tools import ExternalTool


def bwa_index(genome: Path, logdir: Path) -> None:
    """
    Runs bwa index

    Parameters:
        genome (Path): Path to the reference genome to index

    Returns:
        None

    """

    bwa_index = ExternalTool(
        tool="bwa",
        input=f"",
        output=f"",
        params=f"index {genome} ",
        logdir=logdir,
        outfile="",
    )

    ExternalTool.run_tool(bwa_index)


def bwa_paired(
    reads1: Path, reads2: Path, genome: Path, out_sam: Path, threads: int, logdir: Path
) -> None:
    """
    Align paired-end sequencing reads to a reference genome using bwa mem.

    Parameters:
        reads1 (Path): Path to the first set of input FASTQ files containing paired-end reads.
        reads2 (Path): Path to the second set of input FASTQ files containing paired-end reads.
        genome (Path): Path to the reference genome to align the reads against.
        out_sam (Path): Path to the output SAM where the alignment results will be stored.
        threads (int): Number of threads/cores to use for the alignment.
        logdir (Path): Path to the directory where log files will be stored.

    Returns:
        None

    This function uses bwa to perform alignment of paired-end sequencing reads ('reads1' and 'reads2')
    against a reference genome ('genome'). The aligned output is stored in a SAM (Sequence Alignment/Map)
    format file named "temp_bwa.sam" in the 'output' directory.

    The 'threads' parameter specifies the number of CPU cores to use for the alignment, and 'logdir' is the
    directory where log files for the alignment process will be stored.
    """

    bwa_paired = ExternalTool(
        tool="bwa",
        input=f"",
        output=f"",
        params=f"mem -SP -t {str(threads)} {genome} {reads1} {reads2} ",
        logdir=logdir,
        outfile=out_sam,
    )

    ExternalTool.run_tool(bwa_paired, to_stdout=True)


def bwa_single(
    reads1: Path, genome: Path, out_sam: Path, threads: int, logdir: Path
) -> None:
    """
    Align single-end sequencing reads to a reference genome using bwa mem.

    Parameters:
        reads1 (Path): Path to the first set of input FASTQ files containing paired-end reads.
        genome (Path): Path to the reference genome to align the reads against.
        out_sam (Path): Path to the output SAM where the alignment results will be stored.
        threads (int): Number of threads/cores to use for the alignment.
        logdir (Path): Path to the directory where log files will be stored.

    Returns:
        None

    This function uses bwa to perform alignment of paired-end sequencing reads ('reads1' and 'reads2')
    against a reference genome ('genome'). The aligned output is stored in a SAM (Sequence Alignment/Map)
    format file named "temp_bwa.sam" in the 'output' directory.

    The 'threads' parameter specifies the number of CPU cores to use for the alignment, and 'logdir' is the
    directory where log files for the alignment process will be stored.
    """

    bwa_single = ExternalTool(
        tool="bwa",
        input=f"",
        output=f"",
        params=f"mem -SP -t {str(threads)} {genome} {reads1} ",
        logdir=logdir,
        outfile=out_sam,
    )

    ExternalTool.run_tool(bwa_single, to_stdout=True)


def sam_to_bam(sam: Path, bam: Path, threads: int, logdir: Path) -> None:
    """converts sam to bam with samtools
    :param outdir: output directory path
    :param threads: threads
    :param logdir: logdir
    :return:
    """

    sam_to_bam = ExternalTool(
        tool="samtools",
        input="",
        output="",
        params=f" view -h -@ {threads} -b {sam} -o {bam}",
        logdir=logdir,
        outfile="",
    )

    # need to write to stdout
    ExternalTool.run_tool(sam_to_bam, to_stdout=False)


def bam_to_sorted_bam(
    bam: Path, sorted_bam: Path, threads: int, memory_limit: str, logdir: Path
) -> None:
    """converts bam to sorted bam with samtools
    :param outdir: output directory path
    :param threads: threads
    :param logdir: logdir
    :return:
    """

    bam_to_sorted_bam = ExternalTool(
        tool="samtools",
        input="",
        output="",
        params=f" sort -m {memory_limit} -@ {threads} {bam} -o {sorted_bam}",
        logdir=logdir,
        outfile="",
    )

    # need to write to stdout
    ExternalTool.run_tool(bam_to_sorted_bam, to_stdout=False)


def samtools_index(sorted_bam: Path, logdir: Path) -> None:
    """converts bam to sorted bam with samtools
    :param outdir: output directory path
    :param threads: threads
    :param logdir: logdir
    :return:
    """

    samtools_index = ExternalTool(
        tool="samtools",
        input="",
        output="",
        params=f" index -@ {sorted_bam} ",
        logdir=logdir,
        outfile="",
    )

    # need to write to stdout
    ExternalTool.run_tool(samtools_index, to_stdout=False)

src/pypolca/utils/mapping.py
gbouras13/pypolca
import os

import pandas as pd
from Bio import SeqIO
from loguru import logger


def process_blast_output_and_reorient(input, blast_file, out_file, gene: str):
    """Processes

    :param input: input file
    :param blast_file: blast output file
    :param out_file: output file
    :param gene: str - type of gene (dnaA etc)
    :return: blast_success: bool - whether a BLAST hit with a valid start codon was identified
    """

    #  instantiate the returned bool
    blast_success = True

    # define colnames
    col_list = [
        "qseqid",
        "qlen",
        "sseqid",
        "slen",
        "length",
        "qstart",
        "qend",
        "sstart",
        "send",
        "pident",
        "nident",
        "gaps",
        "mismatch",
        "evalue",
        "bitscore",
        "qseq",
        "sseq",
    ]
    # read in the dataframe from BLAST
    try:
        blast_df = pd.read_csv(
            blast_file, delimiter="\t", index_col=False, names=col_list
        )
    except Exception:
        logger.error("There was an issue with parsing the BLAST output file.")

    if isinstance(blast_df, pd.DataFrame) and blast_df.empty:
        logger.info(
            "There were 0 BLAST hits. Please check your input file or try dnaapler custom. If you have assembled an understudied species, this may also be the cause."
        )
        blast_success = False

    # top hit has a start of 1 ########
    # take the top hit - blast sorts by bitscore
    # if the start is 1 of the top hit
    elif blast_df["qstart"][0] == 1:
        logger.info(
            f"Based on the BLAST output top hit, your input is already oriented to begin with {gene}.\n"
            "The input file will be copied to the output so that your pipelines don't break :)"
        )
        # writes to file
        record = SeqIO.read(input, "fasta")
        with open(out_file, "w") as out_fa:
            SeqIO.write(record, out_fa, "fasta")

        blast_success == True

    # top hit
    # prokaryotes can use AUG M, GUG V or UUG L as start codons - e.g. for Pseudomonas aeruginosa PA01  dnaA actually starts with V
    # Therefore, I will require the start codon to be the 1st in the searched sequence match - unlikely to not be but nonetheless
    # Sometimes, the top hit might not be the actual gene of interest (e.g. in phages if the terL is disrupted - like SAOMS1)
    # so in these cases, go down the list and check there is a hit with a legitimate start codon
    # I anticipate this will be rare usually, the first will be it :)
    else:
        gene_found = False
        for i in range(0, len(blast_df.qseq)):
            if blast_df["qseq"][i][0] in ["M", "V", "L"] and (
                blast_df["sstart"][i] == 1
            ):
                reorient_sequence(blast_df, input, out_file, gene, i)
                # warn if the top hit not used
                if i >= 1:
                    logger.warning(
                        f"The top {gene} blastx hit was not chosen to reorient your genome, as it did not begin with a valid start codon."
                    )
                    logger.warning(
                        "A lower hit had a valid start codon which was used instead."
                    )
                gene_found = True
                break
        if gene_found is False:
            logger.info(
                f"{gene} start not identified. Please check your input file or try dnaapler custom. If you have assembled an understudied species, this may also cause this error."
            )
            blast_success = False

    return blast_success


def reorient_sequence(blast_df, input, out_file, gene, i):
    # get the start
    dnaa_start_on_chromosome = blast_df["qstart"][i]
    strand = "forward"
    if blast_df["qstart"][i] > blast_df["qend"][i]:
        strand = "reverse"

    top_hit = blast_df["sseqid"][i]
    top_hit_length = blast_df["slen"][i]
    covered_len = blast_df["length"][i]
    coverage = round(covered_len / top_hit_length * 100, 2)
    ident = blast_df["nident"][i]
    identity = round(ident / covered_len * 100, 2)

    logger.info(
        f"{gene} gene identified. It starts at coordinate {dnaa_start_on_chromosome} on the {strand} strand in your input file."
    )
    logger.info(
        f"The best hit with a valid start codon in the database is {top_hit}, which has length of {top_hit_length} AAs."
    )
    logger.info(
        f"{covered_len} AAs were covered by the best hit, with an overall coverage of {coverage}%."
    )
    logger.info(f"{ident} AAs were identical, with an overall identity of {identity}%.")
    logger.info("Re-orienting.")

    ####################
    # reorientation
    ####################
    record = SeqIO.read(input, "fasta")
    # length of chromosome
    length = len(record.seq)

    # reorient to start at the terminase
    if strand == "forward":
        left = record.seq[(int(dnaa_start_on_chromosome) - 1) : length]
        right = record.seq[0 : int(dnaa_start_on_chromosome) - 1]
        total_seq = left + right

    # revese compliment if the strand is negative
    if strand == "reverse":
        record.seq = record.seq.reverse_complement()
        left = record.seq[(length - int(dnaa_start_on_chromosome)) : length]
        right = record.seq[0 : (length - int(dnaa_start_on_chromosome))]
        total_seq = left + right

    # updates the sequence
    record.seq = total_seq

    # writes to file
    with open(out_file, "w") as out_fa:
        SeqIO.write(record, out_fa, "fasta")


def reorient_sequence_random(input, out_file, start, strand):
    """
    for dnaapler mystery and nearest
    """

    ####################
    # reorientation
    ####################
    record = SeqIO.read(input, "fasta")
    # length of chromosome
    length = len(record.seq)

    # reorient to start at the terminase
    # forward is 1 from pyrodigal
    if strand == 1:
        left = record.seq[(int(start) - 1) : length]
        right = record.seq[0 : int(start) - 1]
        total_seq = left + right

    # revese compliment if the strand is -1 from pyrodigal
    if strand == -1:
        record.seq = record.seq.reverse_complement()
        left = record.seq[(length - int(start)) : length]
        right = record.seq[0 : (length - int(start))]
        total_seq = left + right

    # updates the sequence
    record.seq = total_seq

    # writes to file
    with open(out_file, "w") as out_fa:
        SeqIO.write(record, out_fa, "fasta")


def reorient_single_record_bulk(blast_df, out_file, record, i):
    """
    reorients a single record in dnaapler bulk
    """

    # get the start
    dnaa_start_on_chromosome = blast_df["qstart"][i]
    strand = "forward"
    if blast_df["qstart"][i] > blast_df["qend"][i]:
        strand = "reverse"

    top_hit = blast_df["sseqid"][i]
    top_hit_length = blast_df["slen"][i]
    covered_len = blast_df["length"][i]
    coverage = round(covered_len / top_hit_length * 100, 2)
    ident = blast_df["nident"][i]
    identity = round(ident / covered_len * 100, 2)

    ####################
    # reorientation
    ####################

    # length of chromosome
    length = len(record.seq)

    # reorient to start at the terminase
    if strand == "forward":
        left = record.seq[(int(dnaa_start_on_chromosome) - 1) : length]
        right = record.seq[0 : int(dnaa_start_on_chromosome) - 1]
        total_seq = left + right

    # revese compliment if the strand is negative
    if strand == "reverse":
        record.seq = record.seq.reverse_complement()
        left = record.seq[(length - int(dnaa_start_on_chromosome)) : length]
        right = record.seq[0 : (length - int(dnaa_start_on_chromosome))]
        total_seq = left + right

    # updates the sequence
    record.seq = total_seq

    # writes to file
    with open(out_file, "a") as out_fa:
        SeqIO.write(record, out_fa, "fasta")

    return (
        dnaa_start_on_chromosome,
        strand,
        top_hit,
        top_hit_length,
        covered_len,
        coverage,
        ident,
        identity,
    )


# function to touch create a file
# https://stackoverflow.com/questions/12654772/create-empty-file-using-python
def touch_file(path):
    with open(path, "a"):
        os.utime(path, None)


# to create empty fasta file files
def touch_output_files(out_dir, prefix):
    touch_file(os.path.join(out_dir, prefix + "_reoriented.fasta"))

src/pypolca/utils/processing.py
gbouras13/pypolca
import math
import sys
from pathlib import Path

from Bio import SeqIO


def create_report(vcf: Path, genome: Path, report_file: Path) -> None:
    with open(vcf, "r") as file:
        lines = file.readlines()

    # subs
    numsub = 0
    for line in lines:
        if not line.startswith("#"):
            fields = line.strip().split("\t")
            if len(fields[3]) == 1 and len(fields[4]) == 1:
                parts = fields[9].split(":")
                if int(parts[3]) == 0 and int(parts[5]) > 1:
                    # parts[3] is the number of reads supporting the ref
                    # parts[5] is the number of reads supporting the alternative
                    numsub += 1

    # index

    # Initialize variables
    nind = 0

    # Iterate through the lines and perform the desired operations
    for line in lines:
        # Skip lines starting with #
        if not line.startswith("#"):
            # Split the line using whitespace
            fields = line.split()
            # Check if the fourth and fifth fields have length greater than 1
            if len(fields[3]) > 1 or len(fields[4]) > 1:
                parts = fields[9].split(":")
                if int(parts[3]) == 0 and int(parts[5]) > 1:
                    nind += abs(len(fields[3]) - len(fields[4]))

    # Print the result

    total_size = 0
    for record in SeqIO.parse(genome, "fasta"):
        total_size += len(record.seq)

    numerr = numsub + nind
    erate = numerr / total_size

    qual = 100 - (erate) * 100

    qv = round(-10 * math.log10(erate + 0.0000000001), 2)

    with open(report_file, "w") as f:
        f.write("Stats BEFORE polishing:\n")
        f.write(f"Substitution Errors Found: {numsub}\n")
        f.write(f"Insertion/Deletion Errors Found: {nind}\n")
        f.write(f"Assembly Size: {total_size}\n")
        f.write(f"Consensus Quality Before Polishing: {qual}\n")
        f.write(f"Consensus QV Before Polishing: {qv}\n")

src/pypolca/utils/report.py
gbouras13/pypolca
"""MISC FUNCTIONS
You shouldn't need to tweak these much if at all
"""

import os
import shutil
import subprocess as sp
import sys
import time
from pathlib import Path

import click
from loguru import logger


class OrderedCommands(click.Group):
    """This class will preserve the order of subcommands, which is useful when printing --help"""

    def list_commands(self, ctx: click.Context):
        return list(self.commands)


def pypolca_base(rel_path):
    return os.path.join(os.path.dirname(os.path.realpath(__file__)), rel_path)


def get_version():
    with open(pypolca_base("VERSION"), "r") as f:
        version = f.readline()
    return version


def print_splash():
    click.echo(
        """\b
                         _           
 _ __  _   _ _ __   ___ | | ___ __ _ 
| '_ \| | | | '_ \ / _ \| |/ __/ _` |
| |_) | |_| | |_) | (_) | | (_| (_| |
| .__/ \__, | .__/ \___/|_|\___\__,_|
|_|    |___/|_| 

"""
    )


def echo_click(msg, log=None):
    click.echo(msg, nl=False, err=True)
    if log:
        with open(log, "a") as lo:
            lo.write(msg)


def print_citation():
    with open(pypolca_base("CITATION"), "r") as f:
        for line in f:
            echo_click(line)


log_fmt = (
    "[<green>{time:YYYY-MM-DD HH:mm:ss}</green>] <level>{level: <8}</level> | "
    "<level>{message}</level>"
)

"""
begin and end functions
"""


def begin_pypolca(params) -> None:
    """
    begins pypolca
    params: params a dictionary of params
    returns start time
    """
    # get start time
    start_time = time.time()
    # initial logging stuff
    log_file = os.path.join(params["--output"], f"pypolca_{start_time}.log")
    # adds log file
    logger.add(log_file)
    logger.add(lambda _: sys.exit(1), level="ERROR")
    print_splash()
    logger.info(
        f"pypolca: Standalone Python implementation of the POLCA polisher from MaSuRCA"
    )
    logger.info(f"You are using pypolca version {get_version()}")
    logger.info("Repository homepage is https://github.com/gbouras13/pypolca.")
    logger.info(
        "Written by George Bouras: george.bouras@adelaide.edu.au adapting the original POLCA code by Aleksey Zimin."
    )

    logger.info(f"Listing input parameters.")
    for key, value in params.items():
        logger.info(f"Parameter: {key} {value}.")

    return start_time


def end_pypolca(start_time) -> None:
    """
    finishes pypolca
    """

    # Determine elapsed time
    elapsed_time = time.time() - start_time
    elapsed_time = round(elapsed_time, 2)

    # Show elapsed time for the process
    logger.info("pypolca has finished")
    logger.info("Elapsed time: " + str(elapsed_time) + " seconds")


def copy_file(source_path: Path, destination_path: Path) -> None:
    """
    Copy a file from the source path to the destination path.

    Args:
        source_path (Path): The path to the source file to be copied.
        destination_path (Path): The path to the destination where the file will be copied to.

    Raises:
        FileNotFoundError: If the source file does not exist.
        IsADirectoryError: If the source path is a directory.
        PermissionError: If the user does not have permission to copy the file.
        shutil.Error: If an error occurs during the copy operation.

    Returns:
        None
    """
    # Check if the source file exists and is not a directory
    if not source_path.is_file():
        raise FileNotFoundError(f"Source file '{source_path}' not found.")

    # Use shutil.copy to perform the file copy
    try:
        shutil.copyfile(source_path, destination_path)
    except IsADirectoryError:
        raise IsADirectoryError(
            f"Source path '{source_path}' is a directory, not a file."
        )
    except PermissionError:
        raise PermissionError(
            f"You do not have permission to copy '{source_path}' to '{destination_path}'."
        )
    except shutil.Error as e:
        raise shutil.Error(f"An error occurred while copying the file: {e}")


def remove_directory(dir_path: Path) -> None:
    """
    Remove a directory and its contents if it exists.

    Args:
        dir_path (Path): The Path object representing the directory to be removed.

    Returns:
        None: This function does not return a value.

    Raises:
        FileNotFoundError: If the specified directory does not exist.
        NotADirectoryError: If the specified path is not a directory.

    Example:
        remove_directory(Path('/path/to/directory'))
    """
    if dir_path.exists() and dir_path.is_dir():
        shutil.rmtree(dir_path)

src/pypolca/utils/util.py
gbouras13/pypolca
import gzip
import os
import shutil
import subprocess as sp
import sys
from pathlib import Path

from Bio import SeqIO
from loguru import logger


def instantiate_dirs(output_dir: str, force: bool) -> None:
    """Checks the output directory
    :param out_dir: output directory path
    :param force: force flag
    :param logger: logger
    :return: out_dir: final output directory
    """

    # Checks the output directory
    # remove outdir on force
    logger.add(lambda _: sys.exit(1), level="ERROR")
    logger.info(f"Checking the output directory {output_dir}")
    if force is True:
        if Path(output_dir).exists():
            shutil.rmtree(output_dir)
        else:
            logger.info(
                "--force was specified even though the output directory does not already exist. Continuing."
            )
    else:
        if Path(output_dir).exists():
            logger.error(
                "Output directory already exists and force was not specified. Please specify -f or --force to overwrite the output directory."
            )

    # instantiate outdir
    if Path(output_dir).exists() is False:
        Path(output_dir).mkdir(parents=True, exist_ok=True)


def validate_fasta(input_fasta: Path) -> None:
    """
    Validates  FASTA input - checks that the input is a FASTA
    """
    logger.info(f"Checking that the input file {input_fasta} is in FASTA format.")
    # to get extension
    with open(input_fasta, "r") as handle:
        fasta = SeqIO.parse(handle, "fasta")
        if any(fasta):
            logger.info(f"{input_fasta} file checked.")
        else:
            logger.error(
                f"Error: {input_fasta} file is not in the FASTA format. Please check your input file."
            )


def validate_fastq(file: Path) -> bool:
    """Checks the input fastq is really a fastq
        :param file: fastq file
    :return: zipped - Boolean whether the input fastq is gzipped.
    """

    # to get extension
    filename, file_extension = os.path.splitext(file)
    # flag for whether file is zipped
    zipped = True
    if file_extension == ".gz":
        # if gzipped
        with gzip.open(file, "rt") as handle:
            fastq = SeqIO.parse(handle, "fastq")
            if any(fastq):
                logger.info(f"FASTQ {file} checked")
            else:
                logger.error(f"Input file {file} is not in the FASTQ format.")
    else:
        zipped = False
        with open(file, "r") as handle:
            fastq = SeqIO.parse(handle, "fastq")
            if any(fastq):
                logger.info(f"FASTQ {file} checked")
            else:
                logger.error(f"Input file {file} is not in the FASTQ format.")
    return zipped


def check_memory_limit(memory_limit):
    """
    checks if memory limit is a string that ends with K, G or M.
    """

    memory_limit = memory_limit.upper()
    # -1 gets the last char
    if not any(x in memory_limit[-1] for x in ["K", "G", "M"]):
        logger.error(
            f"Memory limit format must be in format of e.g. 2GB or 2MB (GB being gigabytes of RAM and MB being megabytes of RAM)"
        )


def check_dependencies() -> None:
    """Checks the version of freebayes, bwa and samtools
    :return:
    """

    # samtools
    try:
        process = sp.Popen(["samtools", "--version"], stdout=sp.PIPE, stderr=sp.PIPE)
        samtools_out, _ = process.communicate()
        samtools_out = samtools_out.decode()
        samtools_version = samtools_out.split("\n")[0].split(" ")[1]
        message = f"Samtools v{samtools_version} found."
        logger.info(message)
    except Exception:
        logger.error("Samtools not found.")

    # freebayes
    try:
        process = sp.Popen(["freebayes", "--version"], stdout=sp.PIPE, stderr=sp.PIPE)
        freebayes_out, _ = process.communicate()
        freebayes_out = freebayes_out.decode()
        freebayes_version = freebayes_out.split("\n")[0].split(" ")[2]
        message = f"freebayes {freebayes_version} found."
        logger.info(message)
    except Exception:
        logger.error("freebayes not found.")

    # bwa

    try:
        process = sp.Popen(["bwa"], stdout=sp.PIPE, stderr=sp.PIPE)
        bwa_out, _ = process.communicate()
        # a bit weird
        bwa_out = _.decode()
        version_line = []
        for line in bwa_out.split("\n"):
            if "Version" in line:
                version_line.append(line)
        bwa_version = version_line[0].split(" ")[1]
        message = f"bwa v{bwa_version} found."
        logger.info(message)
    except Exception:
        logger.error("bwa not found")

src/pypolca/utils/validation.py
gbouras13/pypolca
from pathlib import Path

from pypolca.utils.external_tools import ExternalTool


def samtools_faidx(genome: Path, logdir: Path) -> None:
    samtools_faidx = ExternalTool(
        tool="samtools",
        input="",
        output="",
        params=f" faidx {genome} ",
        logdir=logdir,
        outfile="",
    )

    # need to write to stdout
    ExternalTool.run_tool(samtools_faidx, to_stdout=False)


def run_freebayes(
    genome: Path, sorted_bam: Path, output: Path, vcf: Path, logdir: Path
) -> None:
    """
    Run FreeBayes variant calling tool on input data.

    Args:
        genome (Path): Path to the reference genome file.
        sorted_bam (Path): Path to the sorted BAM file containing aligned sequencing reads.
        output (Path): Path to the directory where the output files will be stored.
        vcf (Path): vcf Path
        logdir (Path): Path to the directory where log files will be stored.

    Returns:
        None
    """

    freebayes = ExternalTool(
        tool="freebayes",
        input="",
        output="",
        params=f"  -m 0 --min-coverage 3 -R 0 -p 1 -F 0.2 -E 0 -b {sorted_bam}  -f {genome} -v {vcf} ",
        logdir=logdir,
        outfile="",
    )

    # need to write to stdout
    ExternalTool.run_tool(freebayes, to_stdout=False)

src/pypolca/utils/variants.py
gbouras13/pypolca
"""
Unit tests for dnaapler overall

Usage: pytest .

"""

# import
import os
import shutil

# import functions
import subprocess
import unittest
from pathlib import Path

import pytest

test_data = Path("tests/test_data")


def remove_directory(dir_path):
    if os.path.exists(dir_path):
        shutil.rmtree(dir_path)


@pytest.fixture(scope="session")
def tmp_dir(tmpdir_factory):
    return tmpdir_factory.mktemp("tmp")


# define threads
threads: int = 4


def exec_command(cmnd, stdout=subprocess.PIPE, stderr=subprocess.PIPE):
    """executes shell command and returns stdout if completes exit code 0
    Parameters
    ----------
    cmnd : str
      shell command to be executed
    stdout, stderr : streams
      Default value (PIPE) intercepts process output, setting to None
      blocks this."""

    proc = subprocess.Popen(cmnd, shell=True, stdout=stdout, stderr=stderr)
    out, err = proc.communicate()
    if proc.returncode != 0:
        raise RuntimeError(f"FAILED: {cmnd}\n{err}")
    return out.decode("utf8") if out is not None else None


def test_C347(tmp_dir):
    """test C347"""
    input_fasta: Path = f"{test_data}/C347_flye.fasta"
    r1: Path = f"{test_data}/C347_R1.fastq.gz"
    r2: Path = f"{test_data}/C347_R2.fastq.gz"
    outdir: Path = "output_dir"
    cmd = f"pypolca run -a {input_fasta} -1 {r1} -2 {r2} -t {threads} -o {outdir} -f"
    exec_command(cmd)
    remove_directory(outdir)


def test_C347_single(tmp_dir):
    """test C347 single"""
    input_fasta: Path = f"{test_data}/C347_flye.fasta"
    r1: Path = f"{test_data}/C347_R1.fastq.gz"
    outdir: Path = "output_dir"
    cmd = f"pypolca run -a {input_fasta} -1 {r1} -t {threads} -o {outdir} -f"
    exec_command(cmd)
    remove_directory(outdir)


def test_citation():
    """test citation"""
    cmd = f"pypolca citation"
    exec_command(cmd)


class TestExits(unittest.TestCase):
    """Tests of End to End common failures"""

    def test_FASTQ_as_FASTA(self):
        """test FASTQ as FASTA to fail"""
        with self.assertRaises(RuntimeError):
            input_fasta: Path = f"{test_data}/C347_flye.fasta"
            r1: Path = f"{test_data}/C347_R1.fastq.gz"
            outdir: Path = "output_dir"
            cmd = f"pypolca run -a {r1} -1 {r1} -t {threads} -o {outdir} -f"
            exec_command(cmd)
            remove_directory(outdir)

tests/test_overall.py
gbouras13/pypolca