interleavei

#!/usr/bin/env python

# GY190919

''' interleavei: intelligent FASTA/Q de/interleaving '''


from argparse import ArgumentParser, FileType, SUPPRESS
from signal import signal, SIGPIPE, SIG_DFL
from sys import stderr, stdin, stdout
from typing import Iterator, NamedTuple, TextIO

signal(SIGPIPE, SIG_DFL)  # Gracefully handle downstream PIPE closure


# Classes #####################################################################


class SeqRecord(NamedTuple):
    ''' A class for holding sequence data '''
    seqid: str
    seq: str
    quals: str = ''

    def __len__(self):
        return len(self.seq)

    def print(self, wrap: int=0, fileobj: TextIO=stdout) -> None:
        ''' Print a sequence as FASTA/Q '''
        if self.quals:
            print(f'@{self.seqid}\n{self.seq}\n+\n{self.quals}', file=fileobj)
        elif wrap:
            print(f'>{self.seqid}', file=fileobj)
            for g in grouper(self.seq, wrap):
                print(g, file=fileobj)
        else:
            print(f'>{self.seqid}\n{self.seq}', file=fileobj)


class MultiReader(object):
    ''' A class for reading FASTA/Q records '''

    def __init__(self, inputstream: TextIO, replaceU: bool=False):
        self._fileobj = inputstream
        self._buffer = ' '
        self.replaceU = replaceU

    def __enter__(self):
        return self

    def __exit__(self, type, value, traceback):
        pass

    def __iter__(self):
        return self

    def __next__(self) -> SeqRecord:
        while self._buffer and self._buffer[0] not in ('>', '@'):
            self._buffer = self._fileobj.readline()
        if not self._buffer:
            raise StopIteration
        if self._buffer.startswith('>'):
            seqid = self._buffer[1:].strip('\r\n')
            self._buffer = self._fileobj.readline()
            seqlines = []
            while self._buffer and self._buffer[0] not in ('>', '@'):
                seqlines.append(self._buffer.strip('\r\n'))
                self._buffer = self._fileobj.readline()
            if not (seq := ''.join(seqlines)):
                raise EOFError
            seq = seq.upper()
            if self.replaceU:
                seq = seq.replace('U', 'T')
            return SeqRecord(seqid, seq)
        else:  # self._buffer.startswith('@')
            seqid = self._buffer[1:].strip('\r\n')
            seq = self._fileobj.readline().strip('\r\n')
            self._fileobj.readline()
            quals = self._fileobj.readline().strip('\r\n')
            self._buffer = self._fileobj.readline()
            if not (seq and quals):
                raise EOFError
            seq = seq.upper()
            if self.replaceU:
                seq = seq.replace('U', 'T')
            return SeqRecord(seqid, seq, quals)


# Functions ###################################################################


def grouper(iterable: Iterator, k: int) -> Iterator[Iterator]:
    ''' Yield a memory-mapped `iterable` in chunks of `k`
    !!! If iterable % k != 0 the length of the last chunk will be <k !!! '''
    for i in range(len(iterable) // k + (1 if len(iterable) % k else 0)):
        yield iterable[i*k:i*k+k]


def get_uniform_id(seqid: str) -> str:
    ''' Strip Illumina reads of their pairing information '''
    return seqid.replace('/', ' ').split()[0]


###############################################################################


if __name__ == '__main__':

    parser = ArgumentParser(
        description='intelligent FASTA/Q de/interleaving',
        argument_default=SUPPRESS)
    parser.add_argument(
        'input', nargs='*', type=FileType('r'), default=[stdin],
        help='Unordered input file(s) containing reads to interleave (use "-" or \
        leave blank for stdin)')
    parser.add_argument(
        '-v', '--inverse', nargs=2, type=FileType('w'),
        help='Deinterleave the input into two files')
    parser.add_argument(
        '-w', '--wrap', nargs='?', type=int, const=80, default=0,
        help='Wrap lines at `w` when outputting FASTA (%(const)s when \
        supplied alone)')
    parser.add_argument(
        '-q', '--quiet', action='store_true', default=False,
        help='Quiet! Silence reporting to stderr')
    args = parser.parse_args()

    pairs = {}

    for f in args.input:
        with MultiReader(f) as F:
            for record in F:
                uniform_id = get_uniform_id(record.seqid)
                if pair := pairs.pop(uniform_id, False):
                    if 'inverse' in args:
                        if pair.seqid <= record.seqid:
                            pair.print(args.inverse[0], args.wrap)
                            record.print(args.inverse[1], args.wrap)
                        else:
                            record.print(args.inverse[0], args.wrap)
                            pair.print(args.inverse[1], args.wrap)
                    else:
                        if pair.seqid <= record.seqid:
                            pair.print(wrap=args.wrap)
                            record.print(wrap=args.wrap)
                        else:
                            record.print(wrap=args.wrap)
                            pair.print(wrap=args.wrap)
                else:
                    pairs[uniform_id] = record
        f.close()

    if 'inverse' in args:
        args.inverse[0].close()
        args.inverse[1].close()

    if not args.quiet and len(pairs):
        print(f'{len(pairs)} reads were not paired', file=stderr)