utils.py

import os
import re
import shutil
import tempfile
import base64
from io import StringIO
from Bio import Entrez, SeqIO
from Bio.SeqFeature import SeqFeature, FeatureLocation

Entrez.email = "user@example.com"


def fetch_genbank(accession, out_dir):
    handle = Entrez.efetch(db="nuccore", id=accession, rettype="gbwithparts", retmode="text")
    record = SeqIO.read(handle, "genbank")
    handle.close()
    path = os.path.join(out_dir, f"{accession}.gbk")
    SeqIO.write(record, path, "genbank")
    return path


def fasta_to_gbk(fasta_text, out_dir):
    records = list(SeqIO.parse(StringIO(fasta_text), "fasta"))
    if not records:
        return [], ["Could not parse any sequences from the FASTA input."]

    try:
        import pyrodigal
        orf_finder = pyrodigal.GeneFinder(meta=True)
    except ImportError:
        orf_finder = None

    gbk_files, warnings = [], []

    for record in records:
        record.annotations["molecule_type"] = "DNA"

        if orf_finder is not None:
            try:
                genes = orf_finder.find_genes(str(record.seq).encode())
                for i, gene in enumerate(genes, 1):
                    record.features.append(
                        SeqFeature(
                            FeatureLocation(gene.begin - 1, gene.end, strand=gene.strand),
                            type="CDS",
                            qualifiers={"locus_tag": [f"{record.id}_{i:04d}"]},
                        )
                    )
                if not list(genes):
                    warnings.append(f"{record.id}: no ORFs predicted (sequence may be too short).")
            except Exception as e:
                warnings.append(f"{record.id}: ORF prediction failed — {e}")
        else:
            warnings.append(f"{record.id}: pyrodigal not installed, saving without gene annotation.")

        safe_id = re.sub(r"[^\w\-]", "_", record.id)
        path = os.path.join(out_dir, f"{safe_id}.gbk")
        SeqIO.write(record, path, "genbank")
        gbk_files.append(path)

    return gbk_files, warnings


def _wrap_in_iframe(html_content):
    html_b64 = base64.b64encode(html_content.encode("utf-8")).decode("ascii")
    return (
        f'<iframe src="data:text/html;base64,{html_b64}"'
        ' style="width:100%;height:800px;border:0;" allowfullscreen></iframe>'
    )


def _error_html(msg):
    return f'<p style="color:#c0392b;font-family:sans-serif;padding:1rem">{msg}</p>'


def run_clinker(accessions_text, fasta_text=""):
    accessions = [a.strip() for a in accessions_text.strip().splitlines() if a.strip()]
    has_fasta = bool(fasta_text.strip())

    if not accessions and not has_fasta:
        return _error_html("Please enter at least one accession or a FASTA sequence.")

    work_dir = tempfile.mkdtemp()
    gbk_files, errors = [], []

    for acc in accessions:
        try:
            gbk_files.append(fetch_genbank(acc, work_dir))
        except Exception as e:
            errors.append(f"{acc}: {e}")

    if has_fasta:
        fasta_gbks, fasta_warnings = fasta_to_gbk(fasta_text, work_dir)
        gbk_files.extend(fasta_gbks)
        errors.extend(fasta_warnings)

    if not gbk_files:
        shutil.rmtree(work_dir)
        return _error_html("No usable input files.<br>" + "<br>".join(errors))

    html_path = os.path.join(work_dir, "output.html")
    try:
        try:
            from clinker.main import clinker as clinker_api
            clinker_api(gbk_files, plot=html_path, force=True)
        except ImportError:
            import subprocess
            r = subprocess.run(
                ["clinker", "-f", "-i", "0.225", "-p", html_path] + gbk_files,
                capture_output=True, text=True,
            )
            if r.returncode != 0 and not os.path.exists(html_path):
                raise RuntimeError(r.stderr)
    except Exception as e:
        shutil.rmtree(work_dir)
        return _error_html(f"clinker error: {e}")

    if not os.path.exists(html_path):
        shutil.rmtree(work_dir)
        return _error_html("clinker ran but produced no HTML output.")

    with open(html_path, encoding="utf-8") as f:
        html_content = f.read()

    shutil.rmtree(work_dir)
    return _wrap_in_iframe(html_content)