split3c 0.0.1__py3-none-any.whl

split3c/resite/main.py

@@ -0,0 +1,506 @@
+"""
+This script is a the split3c project, designed to process paired-end FASTQ files by fragmenting DNA sequences at specified restriction enzyme sites.
+
+Copyright © 2024 Samir Bertache
+
+SPDX-License-Identifier: AGPL-3.0-or-later
+
+===============================================================================
+
+This program is free software: you can redistribute it and/or modify it under
+the terms of the GNU Affero General Public License as published by the
+Free Software Foundation, either version 3 of the License, or (at your option)
+any later version.
+
+This program 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 Affero General Public License for more details.
+
+You should have received a copy of the GNU Affero General Public License
+along with this program. If not, see <https://www.gnu.org/licenses/>.
+"""
+
+import argparse
+import logging
+import os
+import sys
+from multiprocessing import Process, Queue
+from typing import Any, List, Optional
+
+from .frag import process_items
+from .pretreatment import partition_threads, search_in_database
+from .read import read_fastq_gzip_simultaneously
+from .write_control import manage_pigz_problems, open_output, write_pairs
+
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version as _version
+
+try:
+    __version__ = _version("parasplit")
+except PackageNotFoundError:
+    __version__ = "0+unknown"
+
+
+def _try_rich() -> Optional[dict[str, Any]]:
+    try:
+        from rich import box
+        from rich.console import Console
+        from rich.panel import Panel
+        from rich.table import Table
+        from rich.theme import Theme
+        from rich.traceback import install
+        from rich_argparse import RichHelpFormatter
+    except Exception:
+        return None
+
+    console = Console(
+        theme=Theme({"info": "dim cyan", "error": "bold red", "warning": "magenta"}),
+        width=100,
+    )
+    install(console=console)
+    return {
+        "console": console,
+        "Panel": Panel,
+        "Table": Table,
+        "box": box,
+        "RichHelpFormatter": RichHelpFormatter,
+    }
+
+
+_R = _try_rich()
+
+
+class MyArgumentParser(argparse.ArgumentParser):
+    def error(self, message: str) -> None:
+        if _R is not None:
+            console = _R["console"]
+            Panel = _R["Panel"]
+            console.print(
+                Panel(
+                    f"[bold red]Error:[/bold red] {message}",
+                    title="Incorrect arguments",
+                    expand=True,
+                    width=100,
+                )
+            )
+            self.print_help()
+            self.exit(2)
+        self.print_usage(sys.stderr)
+        self.exit(2, f"{self.prog}: error: {message}\n")
+
+
+def _formatter_class():
+    if _R is None:
+        return argparse.RawTextHelpFormatter
+
+    from rich_argparse import RawTextRichHelpFormatter
+
+    return lambda prog: RawTextRichHelpFormatter(
+        prog,
+        max_help_position=42,
+        width=110,
+    )
+
+
+def _print_banner() -> None:
+    if _R is None:
+        return
+    console = _R["console"]
+    Panel = _R["Panel"]
+    console.print(
+        Panel(
+            "[bold blue]Parasplit[/bold blue]\n"
+            "Split paired FASTQ at restriction enzyme ligation sites.\n\n"
+            """Features \n
+            Find and Utilize Restriction Enzyme Sites: Automatically identify ligation sites from provided enzyme names and generate regex patterns to locate these sites in sequences.
+            Fragmentation: Split sequences at restriction enzyme sites, creating smaller fragments.
+            Multi-threading: Efficiently handle large datasets by utilizing multiple threads for decompression, fragmentation, and compression.
+            Custom Modes: Supports different pairing modes for sequence fragments.\n\n
+            """
+            "Use --help to see detailed options.",
+            title="[bold green]parasplit[/bold green]",
+            subtitle=f"Version: {__version__}",
+            expand=True,
+            width=100,
+        )
+    )
+    console.print("")
+
+
+def _print_summary(args: argparse.Namespace) -> None:
+    if _R is None:
+        return
+    console = _R["console"]
+    Table = _R["Table"]
+    box = _R["box"]
+
+    t = Table(
+        show_edge=True,
+        title="[bold green]Summary[/bold green]",
+        box=box.HEAVY,
+        width=100,
+    )
+    t.add_column("Key", style="cyan", no_wrap=True)
+    t.add_column("Value", style="magenta")
+
+    t.add_row("--source-forward", str(args.source_forward))
+    t.add_row("--source-reverse", str(args.source_reverse))
+    t.add_row("--output-forward", str(args.output_forward))
+    t.add_row("--output-reverse", str(args.output_reverse))
+    t.add_row("--enzymes", str(args.enzymes))
+    t.add_row("--mode", str(args.mode))
+    t.add_row("--seed-size", str(args.seed_size))
+    t.add_row("--buffer-size", str(args.buffer_size))
+    t.add_row("--num-threads", str(args.num_threads))
+    t.add_row("--borderless", str(args.borderless))
+    t.add_row("--tags", str(args.tags))
+    t.add_row("--force", str(args.force))
+    t.add_row("--verbose", str(args.verbose))
+
+    console.print(t)
+    console.print("")
+
+
+def _parse_enzymes(raw: str) -> list[str]:
+    raw = (raw or "").strip()
+    if not raw or raw == "No restriction enzyme found":
+        return []
+    return [x.strip() for x in raw.split(",") if x.strip()]
+
+
+def validate_args(args: argparse.Namespace) -> None:
+    def _file_exists(path: str, what: str) -> None:
+        if not os.path.exists(path):
+            raise ValueError(f"{what}: file not found: {path}")
+        if not os.path.isfile(path):
+            raise ValueError(f"{what}: not a file: {path}")
+
+    def _parent_writable(path: str, what: str) -> None:
+        parent = os.path.dirname(os.path.abspath(path)) or os.getcwd()
+        if not os.path.exists(parent):
+            raise ValueError(f"{what}: parent directory does not exist: {parent}")
+        if not os.access(parent, os.W_OK):
+            raise ValueError(f"{what}: parent directory not writable: {parent}")
+
+    _file_exists(args.source_forward, "Forward FASTQ")
+    _file_exists(args.source_reverse, "Reverse FASTQ")
+    _parent_writable(args.output_forward, "Output R1")
+    _parent_writable(args.output_reverse, "Output R2")
+
+    if args.output_forward == args.output_reverse:
+        raise ValueError("Output R1 and R2 must be different files")
+
+    if args.num_threads < 5:
+        raise ValueError("--num-threads must be >= 5 (recommended >= 8)")
+
+    if args.seed_size < 0:
+        raise ValueError("--seed-size must be >= 0")
+
+    if args.buffer_size <= 0:
+        raise ValueError("--buffer-size must be >= 1")
+
+    if args.mode not in ("fr", "all", "cover"):
+        raise ValueError("--mode must be one of: fr, all, cover")
+
+    if not args.force:
+        for p in (args.output_forward, args.output_reverse):
+            if os.path.exists(p):
+                raise ValueError(
+                    f"Output already exists: {p} (use --force to overwrite)"
+                )
+
+
+def cut(
+    source_forward: str,
+    source_reverse: str,
+    output_forward: str,
+    output_reverse: str,
+    list_enzyme: List[str],
+    mode,
+    seed_size,
+    buffer_size: int = 100,
+    num_threads: int = 8,
+    borderless: bool = False,
+    tags=None,
+) -> None:
+    """
+    Main function to process sequences based on enzyme restriction sites.
+
+    Parameters:
+        source_forward (str): Input file path for forward reads.
+        source_reverse (str): Input file path for reverse reads.
+        output_forward (str): Output file path for processed forward reads.
+        output_reverse (str): Output file path for processed reverse reads.
+        list_enzyme (List[str]): List of restriction enzymes.
+        mode (str): Mode of pairing fragments, "all" or "fr".
+        seed_size (int): Minimum length of fragments to keep.
+        buffer_size (int, optional): Size of buffer. Defaults to 100.
+        num_threads (int, optional): Number of threads to use for processing. Defaults to 8.
+        borderless (bool, optional): Whether to discard ligation sites (borders). Defaults to False.
+
+    Returns:
+        None
+    """
+    # Threads allocations :
+    TRead, TFrag, TWrite = partition_threads(num_threads)
+
+    # Take the enzyme list and make the ligation site list
+    ligation_site_list = search_in_database(list_enzyme, borderless)
+
+    try:
+        # Input and Output Queues
+        Input_Buffer = Queue(maxsize=2048)
+        Output_buffer = Queue(maxsize=512)
+
+        def read_process():
+            read_fastq_gzip_simultaneously(
+                source_forward, source_reverse, Input_Buffer, TRead, TFrag
+            )
+
+        def process_process_all():
+            process_items(
+                Input_Buffer,
+                Output_buffer,
+                ligation_site_list,
+                seed_size,
+                buffer_size,
+                mode,
+                borderless,
+                tags,
+            )
+
+        def write_process():
+            # IMPORTANT: ouvrir/fermer pigz DANS le process writer
+            outF, outR = open_output(TWrite, output_forward, output_reverse)
+            try:
+                write_pairs(Output_buffer, outF, outR, TFrag)
+            finally:
+                manage_pigz_problems(outF, outR, output_forward, output_reverse)
+
+        # Read fastq files in parallel and asynchronous
+        read_p = Process(target=read_process)
+
+        # Choose mode and Create the executor and dispatch work to it
+        if mode == "all":
+            print("Mode ALL selected")
+            process_p_list = [Process(target=process_process_all) for _ in range(TFrag)]
+        elif mode == "fr":
+            print("Mode FR selected")
+            process_p_list = [Process(target=process_process_all) for _ in range(TFrag)]
+        elif mode == "cover":
+            print("Mode COVER selected")
+            process_p_list = [Process(target=process_process_all) for _ in range(TFrag)]
+        else:
+            print(f"Unknown mode: {mode}")
+            sys.exit(1)
+
+        # Create asynchronous writing
+        write_p = Process(target=write_process)
+
+        # Start processes
+        read_p.start()
+        for p in process_p_list:
+            p.start()
+        write_p.start()
+
+        # Wait for all processes to finish
+        read_p.join()
+        for p in process_p_list:
+            p.join()
+        write_p.join()
+
+        procs = [
+            ("read", read_p),
+            *[(f"proc{i}", p) for i, p in enumerate(process_p_list)],
+            ("write", write_p),
+        ]
+        bad = [(name, p.exitcode) for name, p in procs if p.exitcode not in (0, None)]
+        if bad:
+            raise RuntimeError(f"Subprocess failure(s): {bad}")
+
+    except KeyboardInterrupt:
+        print("Keyboard interrupt detected. Terminating...")
+        sys.exit(0)
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = MyArgumentParser(
+        description=(
+            "Split paired-end FASTQ at restriction enzyme ligation sites.\n"
+            "Example: parasplit -sf R1.fq.gz -sr R2.fq.gz -of out_R1.fq.gz -or out_R2.fq.gz -le DpnII,MboI -m fr -nt 12"
+        ),
+        epilog=(
+            "Examples:\n"
+            "  parasplit -sf R1.fq.gz -sr R2.fq.gz -of out_R1.fq.gz -or out_R2.fq.gz -le DpnII -m fr -nt 12 -sz 20\n"
+            "  parasplit -sf R1.fq.gz -sr R2.fq.gz -of out_R1.fq.gz -or out_R2.fq.gz -le DpnII,MboI -m all -nt 24 --tags o\n"
+        ),
+        formatter_class=_formatter_class(),
+    )
+
+    req = parser.add_argument_group("Inputs")
+    out = parser.add_argument_group("Outputs")
+    perf = parser.add_argument_group("Performance")
+    split = parser.add_argument_group("Split parameters")
+    misc = parser.add_argument_group("Misc")
+
+    req.add_argument(
+        "-sf", "--source_forward", required=True, help="Input FASTQ R1 (gz)."
+    )
+    req.add_argument(
+        "-sr", "--source_reverse", required=True, help="Input FASTQ R2 (gz)."
+    )
+    req.add_argument(
+        "-le",
+        "--list_enzyme",
+        default="",
+        help="Comma-separated restriction enzyme names (e.g. DpnII,MboI). Empty means 'no enzyme'.",
+    )
+
+    out.add_argument(
+        "-of", "--output_forward", required=True, help="Output FASTQ R1 (gz)."
+    )
+    out.add_argument(
+        "-or", "--output_reverse", required=True, help="Output FASTQ R2 (gz)."
+    )
+    out.add_argument(
+        "--force",
+        action="store_true",
+        help="Overwrite output files if they already exist.",
+    )
+
+    perf.add_argument(
+        "-nt", "--num_threads", type=int, default=8, help="Total CPU threads budget."
+    )
+
+    split.add_argument(
+        "-m",
+        "--mode",
+        choices=["fr", "all", "cover"],
+        default="cover",
+        help=(
+            "Pairing mode:\n"
+            "  fr    : one forward fragment + one reverse fragment\n"
+            "  all   : all pairwise fragment combinations\n"
+            "  cover : minimal / near-minimal pairs so every fragment appears at least once (Sufficient to recover all post-processing multiplexe)"
+        ),
+    )
+    split.add_argument(
+        "-sz",
+        "--seed_size",
+        type=int,
+        default=20,
+        help="Minimum fragment length to keep after splitting (0 disables filtering).",
+    )
+    split.add_argument(
+        "--buffer-size",
+        type=int,
+        default=100,
+        help="Chunk size flushed to writer.",
+    )
+    split.add_argument(
+        "-b",
+        "--borderless",
+        action="store_true",
+        help="Discard ligation site borders.",
+    )
+    split.add_argument(
+        "--tags",
+        choices=["origin", "no_annot", "o", "na"],
+        default="o",
+        help=(
+            "Header tagging mode for split reads.\n"
+            "  origin/o : include fragment origin tags (F1,R1,...)\n"
+            "  no_annot/na: keep base name only / No annotation (Not recommended if you wish to filter the pairs after the process.)"
+        ),
+    )
+
+    misc.add_argument(
+        "-v",
+        "--verbose",
+        action="count",
+        default=0,
+        help="Increase verbosity (-v: INFO, -vv: DEBUG).",
+    )
+    misc.add_argument("--version", action="version", version=f"%(prog)s {__version__}")
+
+    return parser
+
+
+def main_cli(argv: Optional[list[str]] = None) -> int:
+    _print_banner()
+
+    parser = build_parser()
+    args = parser.parse_args(argv)
+
+    # logging level
+    level = logging.WARNING
+    if args.verbose == 1:
+        level = logging.INFO
+    elif args.verbose >= 2:
+        level = logging.DEBUG
+    logging.basicConfig(level=level, format="%(levelname)s: %(message)s")
+
+    # normalize enzymes
+    args.enzymes = _parse_enzymes(args.list_enzyme)
+
+    try:
+        validate_args(args)
+    except ValueError as e:
+        if _R is not None:
+            Panel = _R["Panel"]
+            _R["console"].print(
+                Panel(
+                    f"[bold red]{e}[/bold red]",
+                    title="Validation",
+                    expand=True,
+                    width=100,
+                )
+            )
+            parser.print_help()
+        else:
+            print(f"ERROR: {e}", file=sys.stderr)
+            parser.print_help(sys.stderr)
+        return 2
+
+    _print_summary(args)
+
+    try:
+        cut(
+            source_forward=args.source_forward,
+            source_reverse=args.source_reverse,
+            output_forward=args.output_forward,
+            output_reverse=args.output_reverse,
+            list_enzyme=args.enzymes,
+            mode=args.mode,
+            seed_size=args.seed_size,
+            tags=args.tags,
+            buffer_size=args.buffer_size,
+            num_threads=args.num_threads,
+            borderless=args.borderless,
+        )
+    except KeyboardInterrupt:
+        return 130
+    except Exception as e:
+        if _R is not None:
+            Panel = _R["Panel"]
+            _R["console"].print(
+                Panel(
+                    f"[bold red]{e}[/bold red]",
+                    title="Runtime error",
+                    expand=True,
+                    width=100,
+                )
+            )
+        else:
+            print(f"ERROR: {e}", file=sys.stderr)
+        return 1
+
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main_cli())