File size: 21,963 Bytes

d9bc96c

import subprocess
import tempfile
from pathlib import Path
from typing import List, Optional

# This is a placeholder for the MCP decorator.
# In a real MCP environment, this would be provided by the MCP framework.
def tool(func):
    """A dummy decorator to stand in for @mcp.tool()."""
    return func

mcp = type("mcp", (), {"tool": tool})


from mcp.server.fastmcp import FastMCP

SERVER_NAME = 'local_bedtools'
mcp = FastMCP(SERVER_NAME)

@mcp.tool()
def intersect(
    a: Path,
    b: List[Path],
    output: Optional[Path] = None,
    wa: bool = False,
    wb: bool = False,
    loj: bool = False,
    wo: bool = False,
    wao: bool = False,
    u: bool = False,
    c: bool = False,
    v: bool = False,
    ubam: bool = False,
    s: bool = False,
    S: bool = False,
    f: float = 1e-9,
    F: float = 1e-9,
    r: bool = False,
    e: bool = False,
    split: bool = False,
    g: Optional[Path] = None,
    header: bool = False,
    bed: bool = False,
    sorted: bool = False,
    names: Optional[str] = None,
    filenames: bool = False,
    nonamecheck: bool = False,
):
    """
    Find overlapping intervals in two or more BED/GFF/VCF/BAM files.

    This tool allows one to screen for overlaps between two sets of genomic features.
    """
    # Input validation
    if not a.exists():
        raise FileNotFoundError(f"Input file -a does not exist: {a}")
    for b_file in b:
        if not b_file.exists():
            raise FileNotFoundError(f"Input file in -b list does not exist: {b_file}")
    if g and not g.exists():
        raise FileNotFoundError(f"Genome file -g does not exist: {g}")

    exclusive_flags = [u, c, v, wo, wao]
    if sum(exclusive_flags) > 1:
        raise ValueError("Options -u, -c, -v, -wo, -wao are mutually exclusive.")
    if s and S:
        raise ValueError("Options -s and -S are mutually exclusive.")
    if f < 0.0 or f > 1.0:
        raise ValueError("-f (fraction) must be between 0.0 and 1.0.")
    if F < 0.0 or F > 1.0:
        raise ValueError("-F (fraction) must be between 0.0 and 1.0.")
    if sorted and not g:
        raise ValueError("The -sorted option requires a genome file (-g).")

    # Command construction
    cmd = ["bedtools", "intersect", "-a", str(a), "-b"]
    cmd.extend([str(p) for p in b])

    if wa: cmd.append("-wa")
    if wb: cmd.append("-wb")
    if loj: cmd.append("-loj")
    if wo: cmd.append("-wo")
    if wao: cmd.append("-wao")
    if u: cmd.append("-u")
    if c: cmd.append("-c")
    if v: cmd.append("-v")
    if ubam: cmd.append("-ubam")
    if s: cmd.append("-s")
    if S: cmd.append("-S")
    if f != 1e-9: cmd.extend(["-f", str(f)])
    if F != 1e-9: cmd.extend(["-F", str(F)])
    if r: cmd.append("-r")
    if e: cmd.append("-e")
    if split: cmd.append("-split")
    if g: cmd.extend(["-g", str(g)])
    if header: cmd.append("-header")
    if bed: cmd.append("-bed")
    if sorted: cmd.append("-sorted")
    if names: cmd.extend(["-names", names])
    if filenames: cmd.append("-filenames")
    if nonamecheck: cmd.append("-nonamecheck")

    # Subprocess execution
    command_executed = " ".join(cmd)
    output_files = []
    stdout_capture, stderr_capture = "", ""

    try:
        if output:
            output_files.append(str(output))
            with open(output, "w") as f_out:
                result = subprocess.run(
                    cmd, check=True, text=True, stdout=f_out, stderr=subprocess.PIPE
                )
            stderr_capture = result.stderr
        else:
            result = subprocess.run(
                cmd, check=True, text=True, capture_output=True
            )
            stdout_capture = result.stdout
            stderr_capture = result.stderr
    except subprocess.CalledProcessError as e:
        return {
            "command_executed": command_executed,
            "stdout": e.stdout or "",
            "stderr": e.stderr or "CalledProcessError with no stderr.",
            "output_files": [],
            "return_code": e.returncode,
        }

    return {
        "command_executed": command_executed,
        "stdout": stdout_capture,
        "stderr": stderr_capture,
        "output_files": output_files,
    }


@mcp.tool()
def merge(
    i: Path,
    output: Optional[Path] = None,
    s: bool = False,
    S: Optional[str] = None,
    d: int = 0,
    c: Optional[str] = None,
    o: Optional[str] = None,
    header: bool = False,
    delim: str = ";",
    n: bool = False,
    nms: bool = False,
    scores: Optional[str] = None,
    bed: bool = False,
    prec: int = 5,
):
    """
    Merge overlapping features in a BED/GFF/VCF file.

    This tool combines overlapping or "book-ended" features into a single feature.
    """
    # Input validation
    if not i.exists():
        raise FileNotFoundError(f"Input file -i does not exist: {i}")
    if S and S not in ["+", "-"]:
        raise ValueError("Option -S must be either '+' or '-'.")
    if s and S:
        raise ValueError("Options -s and -S are mutually exclusive.")
    if d < 0:
        raise ValueError("Option -d (distance) must be a non-negative integer.")
    if (c and not o) or (o and not c):
        raise ValueError("Options -c and -o must be used together.")
    if c and o:
        if len(c.split(',')) != len(o.split(',')):
            raise ValueError("The number of columns in -c must match the number of operations in -o.")

    # Command construction
    cmd = ["bedtools", "merge", "-i", str(i)]

    if s: cmd.append("-s")
    if S: cmd.extend(["-S", S])
    if d > 0: cmd.extend(["-d", str(d)])
    if c: cmd.extend(["-c", c])
    if o: cmd.extend(["-o", o])
    if header: cmd.append("-header")
    if delim != ";": cmd.extend(["-delim", delim])
    if n: cmd.append("-n")
    if nms: cmd.append("-nms")
    if scores: cmd.extend(["-scores", scores])
    if bed: cmd.append("-bed")
    if prec != 5: cmd.extend(["-prec", str(prec)])

    # Subprocess execution
    command_executed = " ".join(cmd)
    output_files = []
    stdout_capture, stderr_capture = "", ""

    try:
        if output:
            output_files.append(str(output))
            with open(output, "w") as f_out:
                result = subprocess.run(
                    cmd, check=True, text=True, stdout=f_out, stderr=subprocess.PIPE
                )
            stderr_capture = result.stderr
        else:
            result = subprocess.run(
                cmd, check=True, text=True, capture_output=True
            )
            stdout_capture = result.stdout
            stderr_capture = result.stderr
    except subprocess.CalledProcessError as e:
        return {
            "command_executed": command_executed,
            "stdout": e.stdout or "",
            "stderr": e.stderr or "CalledProcessError with no stderr.",
            "output_files": [],
            "return_code": e.returncode,
        }

    return {
        "command_executed": command_executed,
        "stdout": stdout_capture,
        "stderr": stderr_capture,
        "output_files": output_files,
    }


@mcp.tool()
def subtract(
    a: Path,
    b: List[Path],
    output: Optional[Path] = None,
    f: float = 1e-9,
    F: float = 1e-9,
    r: bool = False,
    e: bool = False,
    s: bool = False,
    S: bool = False,
    A: bool = False,
    B: bool = False,
    N: bool = False,
    header: bool = False,
    g: Optional[Path] = None,
):
    """
    Remove overlapping intervals from a BED/GFF/VCF file.

    This tool removes portions of features in file A that are overlapped by features in file(s) B.
    """
    # Input validation
    if not a.exists():
        raise FileNotFoundError(f"Input file -a does not exist: {a}")
    for b_file in b:
        if not b_file.exists():
            raise FileNotFoundError(f"Input file in -b list does not exist: {b_file}")
    if g and not g.exists():
        raise FileNotFoundError(f"Genome file -g does not exist: {g}")
    if A and B:
        raise ValueError("Options -A and -B are mutually exclusive.")
    if s and S:
        raise ValueError("Options -s and -S are mutually exclusive.")

    # Command construction
    cmd = ["bedtools", "subtract", "-a", str(a), "-b"]
    cmd.extend([str(p) for p in b])

    if f != 1e-9: cmd.extend(["-f", str(f)])
    if F != 1e-9: cmd.extend(["-F", str(F)])
    if r: cmd.append("-r")
    if e: cmd.append("-e")
    if s: cmd.append("-s")
    if S: cmd.append("-S")
    if A: cmd.append("-A")
    if B: cmd.append("-B")
    if N: cmd.append("-N")
    if header: cmd.append("-header")
    if g: cmd.extend(["-g", str(g)])

    # Subprocess execution
    command_executed = " ".join(cmd)
    output_files = []
    stdout_capture, stderr_capture = "", ""

    try:
        if output:
            output_files.append(str(output))
            with open(output, "w") as f_out:
                result = subprocess.run(
                    cmd, check=True, text=True, stdout=f_out, stderr=subprocess.PIPE
                )
            stderr_capture = result.stderr
        else:
            result = subprocess.run(
                cmd, check=True, text=True, capture_output=True
            )
            stdout_capture = result.stdout
            stderr_capture = result.stderr
    except subprocess.CalledProcessError as e:
        return {
            "command_executed": command_executed,
            "stdout": e.stdout or "",
            "stderr": e.stderr or "CalledProcessError with no stderr.",
            "output_files": [],
            "return_code": e.returncode,
        }

    return {
        "command_executed": command_executed,
        "stdout": stdout_capture,
        "stderr": stderr_capture,
        "output_files": output_files,
    }


@mcp.tool()
def slop(
    i: Path,
    g: Path,
    output: Optional[Path] = None,
    b: int = 0,
    l: int = 0,
    r: int = 0,
    s: bool = False,
    pct: bool = False,
    header: bool = False,
):
    """
    Increase the size of features in a BED/GFF/VCF file.

    This tool will increase the size of each feature in a feature file by a user-defined number of bases.
    """
    # Input validation
    if not i.exists():
        raise FileNotFoundError(f"Input file -i does not exist: {i}")
    if not g.exists():
        raise FileNotFoundError(f"Genome file -g does not exist: {g}")
    if b != 0 and (l != 0 or r != 0):
        raise ValueError("Option -b cannot be used with -l or -r.")

    # Command construction
    cmd = ["bedtools", "slop", "-i", str(i), "-g", str(g)]

    if b != 0: cmd.extend(["-b", str(b)])
    if l != 0: cmd.extend(["-l", str(l)])
    if r != 0: cmd.extend(["-r", str(r)])
    if s: cmd.append("-s")
    if pct: cmd.append("-pct")
    if header: cmd.append("-header")

    # Subprocess execution
    command_executed = " ".join(cmd)
    output_files = []
    stdout_capture, stderr_capture = "", ""

    try:
        if output:
            output_files.append(str(output))
            with open(output, "w") as f_out:
                result = subprocess.run(
                    cmd, check=True, text=True, stdout=f_out, stderr=subprocess.PIPE
                )
            stderr_capture = result.stderr
        else:
            result = subprocess.run(
                cmd, check=True, text=True, capture_output=True
            )
            stdout_capture = result.stdout
            stderr_capture = result.stderr
    except subprocess.CalledProcessError as e:
        return {
            "command_executed": command_executed,
            "stdout": e.stdout or "",
            "stderr": e.stderr or "CalledProcessError with no stderr.",
            "output_files": [],
            "return_code": e.returncode,
        }

    return {
        "command_executed": command_executed,
        "stdout": stdout_capture,
        "stderr": stderr_capture,
        "output_files": output_files,
    }


@mcp.tool()
def getfasta(
    fi: Path,
    bed: Path,
    fo: Optional[Path] = None,
    s: bool = False,
    split: bool = False,
    name: bool = False,
    name_plus: bool = False,
    tab: bool = False,
    full_header: bool = False,
):
    """
    Extract DNA sequences from a FASTA file based on BED/GFF/VCF coordinates.
    """
    # Input validation
    if not fi.exists():
        raise FileNotFoundError(f"FASTA input file -fi does not exist: {fi}")
    if not bed.exists():
        raise FileNotFoundError(f"BED/GFF/VCF file -bed does not exist: {bed}")
    if name and name_plus:
        raise ValueError("Options -name and -name+ are mutually exclusive.")

    # Command construction
    cmd = ["bedtools", "getfasta", "-fi", str(fi), "-bed", str(bed)]

    if s: cmd.append("-s")
    if split: cmd.append("-split")
    if name: cmd.append("-name")
    if name_plus: cmd.append("-name+")
    if tab: cmd.append("-tab")
    if full_header: cmd.append("-fullHeader")
    if fo: cmd.extend(["-fo", str(fo)])

    # Subprocess execution
    command_executed = " ".join(cmd)
    output_files = []
    stdout_capture, stderr_capture = "", ""

    try:
        # getfasta is unique in that it has a dedicated -fo parameter
        # and doesn't typically write to stdout if -fo is used.
        result = subprocess.run(
            cmd, check=True, text=True, capture_output=True
        )
        stdout_capture = result.stdout
        stderr_capture = result.stderr
        if fo:
            output_files.append(str(fo))

    except subprocess.CalledProcessError as e:
        return {
            "command_executed": command_executed,
            "stdout": e.stdout or "",
            "stderr": e.stderr or "CalledProcessError with no stderr.",
            "output_files": [],
            "return_code": e.returncode,
        }

    return {
        "command_executed": command_executed,
        "stdout": stdout_capture,
        "stderr": stderr_capture,
        "output_files": output_files,
    }


@mcp.tool()
def sort(
    i: Path,
    output: Optional[Path] = None,
    header: bool = False,
    g: Optional[Path] = None,
    faidx: Optional[Path] = None,
    sizeA: bool = False,
    sizeD: bool = False,
    chrThenSizeA: bool = False,
    chrThenSizeD: bool = False,
    chrThenScoreA: bool = False,
    chrThenScoreD: bool = False,
):
    """
    Sort a BED/GFF/VCF file by chromosome and then by start position.
    """
    # Input validation
    if not i.exists():
        raise FileNotFoundError(f"Input file -i does not exist: {i}")
    if g and not g.exists():
        raise FileNotFoundError(f"Genome file -g does not exist: {g}")
    if faidx and not faidx.exists():
        raise FileNotFoundError(f"FASTA index file -faidx does not exist: {faidx}")

    sort_flags = [sizeA, sizeD, chrThenSizeA, chrThenSizeD, chrThenScoreA, chrThenScoreD]
    if sum(sort_flags) > 1:
        raise ValueError("Multiple sorting order flags were provided, but only one is allowed.")

    # Command construction
    cmd = ["bedtools", "sort", "-i", str(i)]

    if header: cmd.append("-header")
    if g: cmd.extend(["-g", str(g)])
    if faidx: cmd.extend(["-faidx", str(faidx)])
    if sizeA: cmd.append("-sizeA")
    if sizeD: cmd.append("-sizeD")
    if chrThenSizeA: cmd.append("-chrThenSizeA")
    if chrThenSizeD: cmd.append("-chrThenSizeD")
    if chrThenScoreA: cmd.append("-chrThenScoreA")
    if chrThenScoreD: cmd.append("-chrThenScoreD")

    # Subprocess execution
    command_executed = " ".join(cmd)
    output_files = []
    stdout_capture, stderr_capture = "", ""

    try:
        if output:
            output_files.append(str(output))
            with open(output, "w") as f_out:
                result = subprocess.run(
                    cmd, check=True, text=True, stdout=f_out, stderr=subprocess.PIPE
                )
            stderr_capture = result.stderr
        else:
            result = subprocess.run(
                cmd, check=True, text=True, capture_output=True
            )
            stdout_capture = result.stdout
            stderr_capture = result.stderr
    except subprocess.CalledProcessError as e:
        return {
            "command_executed": command_executed,
            "stdout": e.stdout or "",
            "stderr": e.stderr or "CalledProcessError with no stderr.",
            "output_files": [],
            "return_code": e.returncode,
        }

    return {
        "command_executed": command_executed,
        "stdout": stdout_capture,
        "stderr": stderr_capture,
        "output_files": output_files,
    }


@mcp.tool()
def coverage(
    a: Path,
    b: List[Path],
    output: Optional[Path] = None,
    s: bool = False,
    S: bool = False,
    f: float = 1e-9,
    F: float = 1e-9,
    r: bool = False,
    e: bool = False,
    split: bool = False,
    d: bool = False,
    counts: bool = False,
    hist: bool = False,
    mean: bool = False,
    g: Optional[Path] = None,
    header: bool = False,
    sorted: bool = False,
):
    """
    Compute the coverage of features in file A on features in file(s) B.
    """
    # Input validation
    if not a.exists():
        raise FileNotFoundError(f"Input file -a does not exist: {a}")
    for b_file in b:
        if not b_file.exists():
            raise FileNotFoundError(f"Input file in -b list does not exist: {b_file}")
    if g and not g.exists():
        raise FileNotFoundError(f"Genome file -g does not exist: {g}")
    if s and S:
        raise ValueError("Options -s and -S are mutually exclusive.")
    if sorted and not g:
        raise ValueError("The -sorted option requires a genome file (-g).")

    # Command construction
    cmd = ["bedtools", "coverage", "-a", str(a), "-b"]
    cmd.extend([str(p) for p in b])

    if s: cmd.append("-s")
    if S: cmd.append("-S")
    if f != 1e-9: cmd.extend(["-f", str(f)])
    if F != 1e-9: cmd.extend(["-F", str(F)])
    if r: cmd.append("-r")
    if e: cmd.append("-e")
    if split: cmd.append("-split")
    if d: cmd.append("-d")
    if counts: cmd.append("-counts")
    if hist: cmd.append("-hist")
    if mean: cmd.append("-mean")
    if g: cmd.extend(["-g", str(g)])
    if header: cmd.append("-header")
    if sorted: cmd.append("-sorted")

    # Subprocess execution
    command_executed = " ".join(cmd)
    output_files = []
    stdout_capture, stderr_capture = "", ""

    try:
        if output:
            output_files.append(str(output))
            with open(output, "w") as f_out:
                result = subprocess.run(
                    cmd, check=True, text=True, stdout=f_out, stderr=subprocess.PIPE
                )
            stderr_capture = result.stderr
        else:
            result = subprocess.run(
                cmd, check=True, text=True, capture_output=True
            )
            stdout_capture = result.stdout
            stderr_capture = result.stderr
    except subprocess.CalledProcessError as e:
        return {
            "command_executed": command_executed,
            "stdout": e.stdout or "",
            "stderr": e.stderr or "CalledProcessError with no stderr.",
            "output_files": [],
            "return_code": e.returncode,
        }

    return {
        "command_executed": command_executed,
        "stdout": stdout_capture,
        "stderr": stderr_capture,
        "output_files": output_files,
    }


@mcp.tool()
def genomecov(
    i: Path,
    g: Path,
    output: Optional[Path] = None,
    d: bool = False,
    dz: bool = False,
    bga: bool = False,
    bg: bool = False,
    scale: float = 1.0,
    pc: bool = False,
    fs: bool = False,
    split: bool = False,
    strand: Optional[str] = None,
    max: int = 0,
    trackline: bool = False,
    trackopts: Optional[str] = None,
):
    """
    Compute genome-wide coverage of a feature file.
    """
    # Input validation
    if not i.exists():
        raise FileNotFoundError(f"Input file -i does not exist: {i}")
    if not g.exists():
        raise FileNotFoundError(f"Genome file -g does not exist: {g}")
    if strand and strand not in ["+", "-"]:
        raise ValueError("Option -strand must be either '+' or '-'.")
    if max < 0:
        raise ValueError("Option -max must be a non-negative integer.")
    output_modes = [d, dz, bga, bg]
    if sum(output_modes) > 1:
        raise ValueError("Output format flags (-d, -dz, -bga, -bg) are mutually exclusive.")

    # Command construction
    cmd = ["bedtools", "genomecov", "-i", str(i), "-g", str(g)]

    if d: cmd.append("-d")
    if dz: cmd.append("-dz")
    if bga: cmd.append("-bga")
    if bg: cmd.append("-bg")
    if scale != 1.0: cmd.extend(["-scale", str(scale)])
    if pc: cmd.append("-pc")
    if fs: cmd.append("-fs")
    if split: cmd.append("-split")
    if strand: cmd.extend(["-strand", strand])
    if max > 0: cmd.extend(["-max", str(max)])
    if trackline: cmd.append("-trackline")
    if trackopts: cmd.extend(["-trackopts", trackopts])

    # Subprocess execution
    command_executed = " ".join(cmd)
    output_files = []
    stdout_capture, stderr_capture = "", ""

    try:
        if output:
            output_files.append(str(output))
            with open(output, "w") as f_out:
                result = subprocess.run(
                    cmd, check=True, text=True, stdout=f_out, stderr=subprocess.PIPE
                )
            stderr_capture = result.stderr
        else:
            result = subprocess.run(
                cmd, check=True, text=True, capture_output=True
            )
            stdout_capture = result.stdout
            stderr_capture = result.stderr
    except subprocess.CalledProcessError as e:
        return {
            "command_executed": command_executed,
            "stdout": e.stdout or "",
            "stderr": e.stderr or "CalledProcessError with no stderr.",
            "output_files": [],
            "return_code": e.returncode,
        }

    return {
        "command_executed": command_executed,
        "stdout": stdout_capture,
        "stderr": stderr_capture,
        "output_files": output_files,
    }

if __name__ == "__main__":
    mcp.run(transport="stdio")