split3c 0.0.1__py3-none-any.whl

split3c/nssite/bam.py

@@ -0,0 +1,299 @@
+def get_bam_headers(bam_for_path, bam_rev_path):
+    """
+    Ouvre les fichiers BAM, extrait leurs headers sous forme de dictionnaire,
+    et referme les fichiers immédiatement.
+
+    Returns:
+        tuple: (header_dict_forward, header_dict_reverse)
+    """
+    import pysam
+
+    # Lecture du header Forward
+    with pysam.AlignmentFile(bam_for_path, "rb") as bf:
+        header_for = bf.header.to_dict()
+
+    # Lecture du header Reverse
+    with pysam.AlignmentFile(bam_rev_path, "rb") as br:
+        header_rev = br.header.to_dict()
+
+    return header_for, header_rev
+
+
+def read_bam_pair(
+    bam_for_file,
+    bam_rev_file,
+    input_queue,
+    num_processes,
+    bam_threads=1,
+    batch_size=500,
+):
+    """
+    Read simultaneously two BAM files and put read pairs into an input queue.
+
+    Parameters:
+        bam_for_file (str): Path to the forward BAM file.
+        bam_rev_file (str): Path to the reverse BAM file.
+        Input_Queue (Queue): Queue to store read pairs.
+        TFrag (int): Number of fragmenting threads.
+    """
+    import sys
+
+    import pysam
+
+    batch = []
+    try:
+        with (
+            pysam.AlignmentFile(bam_for_file, "rb", threads=bam_threads) as bam_for,
+            pysam.AlignmentFile(bam_rev_file, "rb", threads=bam_threads) as bam_rev,
+        ):
+            for read_for, read_rev in zip(bam_for, bam_rev):
+                if read_for and read_rev:
+                    # Convert read objects to serializable format
+                    batch.append((read_for.to_string(), read_rev.to_string()))
+
+                if len(batch) >= batch_size:
+                    input_queue.put(batch)
+                    batch = []
+
+        if batch:
+            input_queue.put(batch)
+
+    except Exception as e:
+        print(f"Error: with {bam_for_file} or {bam_rev_file}, {e}")
+        sys.exit(1)
+    finally:
+        for _ in range(num_processes):
+            input_queue.put(None)
+
+
+def write_bam_pair_from_sam(
+    queue,
+    out_bam_f_path,
+    out_bam_r_path,
+    header_for_dict,
+    header_rev_dict,
+    num_procs_finished_signal,
+    bam_threads=1,
+):
+    """
+    Écrit les paires BAM en utilisant les dictionnaires de header fournis.
+    """
+    import sys
+
+    import pysam
+
+    try:
+        header_for = pysam.AlignmentHeader.from_dict(header_for_dict)
+        header_rev = pysam.AlignmentHeader.from_dict(header_rev_dict)
+
+        with (
+            pysam.AlignmentFile(
+                out_bam_f_path, "wb", header=header_for, threads=bam_threads
+            ) as out_f,
+            pysam.AlignmentFile(
+                out_bam_r_path, "wb", header=header_rev, threads=bam_threads
+            ) as out_r,
+        ):
+            finished = 0
+            while finished < num_procs_finished_signal:
+                batch = queue.get()
+
+                if batch is None:
+                    finished += 1
+                    continue
+
+                # batch = list[(sam_f, sam_r)]
+                for sam_f, sam_r in batch:
+                    read_f = pysam.AlignedSegment.fromstring(sam_f, header_for)
+                    read_r = pysam.AlignedSegment.fromstring(sam_r, header_rev)
+
+                    out_f.write(read_f)
+                    out_r.write(read_r)
+
+    except Exception as e:
+        print(f"Error in write_bam_pair_from_sam: {e}", file=sys.stderr)
+        sys.exit(1)
+
+
+#########
+# One bam option
+########
+
+
+def get_bam_header_single(bam_path):
+    """
+    Ouvre un BAM unique, extrait son header sous forme de dictionnaire,
+    puis referme le fichier.
+
+    Returns
+    -------
+    dict
+        Header BAM sous forme de dictionnaire.
+    """
+    import pysam
+
+    with pysam.AlignmentFile(bam_path, "rb") as bam:
+        return bam.header.to_dict()
+
+
+def _pair_reads_from_single_bam(read_a, read_b, strict=True):
+    """
+    Ordonne deux lectures provenant d'un BAM interleavé en (forward/read1, reverse/read2).
+
+    Paramètres
+    ----------
+    read_a, read_b : pysam.AlignedSegment
+    strict : bool
+        Si True, lève une erreur en cas d'incohérence forte.
+        Si False, tente un fallback par ordre d'apparition.
+
+    Returns
+    -------
+    tuple
+        (read_for, read_rev)
+
+    Notes
+    -----
+    On utilise en priorité les flags 0x40 / 0x80 (is_read1 / is_read2).
+    """
+    if read_a is None or read_b is None:
+        raise ValueError("Pairing failure: one of the reads is None.")
+
+    if read_a.query_name != read_b.query_name:
+        raise ValueError(
+            f"Interleaved BAM not synchronized: "
+            f"{read_a.query_name!r} != {read_b.query_name!r}"
+        )
+
+    # Fallback permissif: ordre d'apparition
+    return read_a, read_b
+
+
+def read_bam_interleaved(
+    bam_file,
+    input_queue,
+    num_processes,
+    bam_threads=1,
+    batch_size=500,
+    strict=True,
+):
+    """
+    Lit un BAM unique interleavé (une ligne forward/read1 suivie de la ligne reverse/read2) et envoie des batchs de paires SAM dans input_queue.
+
+    Contrat de sortie identique à read_bam_pair:
+        batch = list[(sam_f, sam_r)]
+
+    Paramètres
+    ----------
+    bam_file : str
+        Chemin vers un BAM unique interleavé.
+    input_queue : multiprocessing.Queue
+    num_processes : int
+        Nombre de workers compute, pour envoyer les sentinelles None.
+    bam_threads : int
+        Threads pysam/htslib.
+    batch_size : int
+        Taille des batchs.
+    strict : bool
+        Si True, échoue si les paires ne sont pas parfaitement cohérentes.
+
+    Exigences
+    ---------
+    Le BAM doit être ordonné par nom ou au minimum avoir les deux mates consécutives.
+    """
+    import sys
+
+    import pysam
+
+    batch = []
+    pending = None
+
+    try:
+        with pysam.AlignmentFile(bam_file, "rb", threads=bam_threads) as bam:
+            for read in bam:
+                if pending is None:
+                    pending = read
+                    continue
+
+                read_for, read_rev = _pair_reads_from_single_bam(
+                    pending, read, strict=strict
+                )
+                batch.append((read_for.to_string(), read_rev.to_string()))
+                pending = None
+
+                if len(batch) >= batch_size:
+                    input_queue.put(batch)
+                    batch = []
+
+        if pending is not None:
+            raise ValueError(
+                f"Odd number of records in interleaved BAM or dangling read: "
+                f"{pending.query_name!r}"
+            )
+
+        if batch:
+            input_queue.put(batch)
+
+    except Exception as e:
+        print(f"Error with interleaved BAM {bam_file}: {e}", file=sys.stderr)
+        sys.exit(1)
+    finally:
+        for _ in range(num_processes):
+            input_queue.put(None)
+
+
+def write_bam_interleaved_from_sam(
+    queue,
+    out_bam_path,
+    header_dict,
+    num_procs_finished_signal,
+    bam_threads=1,
+):
+    """
+    Écrit les paires BAM non splittables dans un BAM unique interleavé.
+
+    Contrat d'entrée:
+        queue contient des batchs list[(sam_f, sam_r)]
+
+    Sortie:
+        un seul BAM avec read1 puis read2 à la suite.
+
+    Paramètres
+    ----------
+    queue : multiprocessing.Queue
+    out_bam_path : str
+    header_dict : dict
+    num_procs_finished_signal : int
+    bam_threads : int
+    """
+    import sys
+
+    import pysam
+
+    try:
+        header = pysam.AlignmentHeader.from_dict(header_dict)
+
+        with pysam.AlignmentFile(
+            out_bam_path,
+            "wb",
+            header=header,
+            threads=bam_threads,
+        ) as out_bam:
+            finished = 0
+            while finished < num_procs_finished_signal:
+                batch = queue.get()
+
+                if batch is None:
+                    finished += 1
+                    continue
+
+                for sam_f, sam_r in batch:
+                    read_f = pysam.AlignedSegment.fromstring(sam_f, header)
+                    read_r = pysam.AlignedSegment.fromstring(sam_r, header)
+
+                    out_bam.write(read_f)
+                    out_bam.write(read_r)
+
+    except Exception as e:
+        print(f"Error in write_bam_interleaved_from_sam: {e}", file=sys.stderr)
+        sys.exit(1)

split3c/nssite/fastq.py

@@ -0,0 +1,148 @@
+def open_output(output_forward, output_reverse, write_processes):
+    """
+    Open output files for writing with pigz compression.
+
+    Parameters:
+        write_processes (int): Number of threads for writing.
+        output_forward (str): Path to the forward output file.
+        output_reverse (str): Path to the reverse output file.
+
+    Returns:
+        out_f (subprocess.Popen): Process for the forward output.
+        out_r (subprocess.Popen): Process for the reverse output.
+    """
+    import signal
+    import subprocess
+
+    from .auxiliary import signal_handler
+
+    # Open output files for writing
+    out_f = subprocess.Popen(
+        args=["pigz", "-c", "-p", str(write_processes)],
+        stdin=subprocess.PIPE,
+        stdout=open(file=output_forward, mode="wb"),
+    )
+    out_r = subprocess.Popen(
+        ["pigz", "-c", "-p", str(write_processes)],
+        stdin=subprocess.PIPE,
+        stdout=open(file=output_reverse, mode="wb"),
+    )
+
+    # Register signal handlers
+    signal.signal(
+        signal.SIGINT,
+        lambda sig, frame: signal_handler(
+            sig=sig, frame=frame, out_f=out_f, out_r=out_r
+        ),
+    )
+    signal.signal(
+        signal.SIGTSTP,
+        lambda sig, frame: signal_handler(
+            sig=sig, frame=frame, out_f=out_f, out_r=out_r
+        ),
+    )
+    return out_f, out_r
+
+
+def write_fastq_pair(
+    output_queue, output_forward, output_reverse, num_process, write_processes
+):
+    """
+    Write FastQ file pairs to the output using data from the output queue.
+
+    Parameters:
+        output_queue (Queue): Queue to get processed read pairs.
+        out_f (subprocess.Popen): Process for the forward output.
+        out_r (subprocess.Popen): Process for the reverse output.
+        num_process (int): Number of fragmenting threads.
+    """
+    import sys
+
+    out_f, out_r = open_output(output_forward, output_reverse, write_processes)
+    while num_process > 0:
+        try:
+            data = output_queue.get()
+            if data is None:
+                num_process -= 1
+                continue
+
+            f_block = data[0]
+            r_block = data[1]
+
+            if out_f.stdin is None or out_r.stdin is None:
+                raise ValueError("pigz stdin closed")
+
+            out_f.stdin.write(f_block.encode("utf-8"))
+            out_r.stdin.write(r_block.encode("utf-8"))
+        except Exception as e:
+            print(f"Error in write_pairs: {e}")
+            manage_pigz_errors(out_f, out_r, output_forward, output_reverse)
+            sys.exit(1)
+    ensure_ending(out_f, out_r)
+
+
+def ensure_ending(out_f, out_r):
+    """
+    Examples
+    --------
+    >>> class _P:
+    ...     def __init__(self):
+    ...         self.stdin = self
+    ...         self.closed = False
+    ...         self.waited = False
+    ...         self.terminated = False
+    ...     def close(self):
+    ...         self.closed = True
+    ...     def wait(self):
+    ...         self.waited = True
+    ...     def terminate(self):
+    ...         self.terminated = True
+    ...
+    >>> pf, pr = _P(), _P()
+    >>> ensure_ending(pf, pr)
+    >>> (pf.closed, pr.closed, pf.waited, pr.waited, pf.terminated, pr.terminated)
+    (True, True, True, True, True, True)
+    """
+    if out_f.stdin is not None:
+        out_f.stdin.close()
+    if out_r.stdin is not None:
+        out_r.stdin.close()
+    out_f.wait()
+    out_r.wait()
+    out_f.terminate()
+    out_r.terminate()
+
+
+def manage_pigz_errors(out_f, out_r, output_forward, output_reverse):
+    """
+    Manage pigz process termination and check for errors.
+
+    Examples
+    --------
+    >>> class _P:
+    ...     def __init__(self, rc): self.stdin=None; self.returncode=rc; self.waited=False; self.terminated=False
+    ...     def wait(self): self.waited=True
+    ...     def terminate(self): self.terminated=True
+    ...
+    >>> pf, pr = _P(0), _P(0)
+    >>> manage_pigz_errors(pf, pr, 'F.fq.gz', 'R.fq.gz')
+    >>> pf.waited and pr.waited and pf.terminated and pr.terminated
+    True
+    >>> pf, pr = _P(1), _P(2)
+    >>> manage_pigz_errors(pf, pr, 'F.fq.gz', 'R.fq.gz')  # doctest: +ELLIPSIS
+    Error in pigz command for file F.fq.gz
+    Error in pigz command for file R.fq.gz
+    """
+    if out_f.stdin is not None:
+        out_f.stdin.close()
+    if out_r.stdin is not None:
+        out_r.stdin.close()
+    out_f.wait()
+    out_r.wait()
+
+    if out_f.returncode != 0:
+        print(f"Error in pigz command for file {output_forward}")
+    if out_r.returncode != 0:
+        print(f"Error in pigz command for file {output_reverse}")
+    out_f.terminate()
+    out_r.terminate()