linearise

#!/usr/bin/env python

# GY190919

''' linearise: FASTA <-> TSV conversion '''


from argparse import ArgumentParser, FileType, SUPPRESS
from signal import signal, SIGPIPE, SIG_DFL
from sys import 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 printseq(seqid: str, seq: str, quals: str='', wrap: int=0, fileobj: TextIO=stdout) -> None:
    ''' Appropraitely prints a sequence '''
    if quals:
        print(f'@{seqid}\n{seq}\n+\n{quals}', file=fileobj)
    elif wrap:
        print(f'>{seqid}', file=fileobj)
        for chunk in grouper(seq, wrap):
            print(chunk, file=fileobj)
    else:
        print(f'>{seqid}\n{seq}', file=fileobj)


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


if __name__ == '__main__':

    parser = ArgumentParser(
        description='FASTA/Q <-> TSV conversion',
        epilog='Converts to and from a standardised 3 column \
        (ID \\t Seq \\t Qual) format. Linearised FASTA data will have no \
        entry in the third column and both two-column data and data where the \
        third column is empty will be re-interpreted as FASTA with `-v`.',
        argument_default=SUPPRESS)
    parser.add_argument(
        'input', nargs='*', type=FileType('r'), default=[stdin],
        help='Input file(s) (use "-" or leave blank for stdin)')
    parser.add_argument(
        '-v', '--inverse', action='store_true',
        help='Convert delimited lines back to FASTA/Q')
    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)')
    args = parser.parse_args()

    if 'inverse' not in args:
        for f in args.input:
            with MultiReader(f) as F:
                for record in F:
                    print('\t'.join(record))
            f.close()
    else:
        for f in args.input:
            for l in f:
                printseq(*l.strip('\r\n').split('\t'), wrap=args.wrap)
            f.close()