rc

#!/usr/bin/env python

# GY190923

''' rc: reverse complement '''


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


# Globals #####################################################################


NUCL_COMP = str.maketrans('ATUCGXNRYWSKMDVHB', 'TAAGCXNYRWSMKHBDV')


# 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 rc(record: SeqRecord) -> SeqRecord:
    ''' Reverse complement a SeqRecord '''
    return SeqRecord(
        record.seqid,
        record.seq[::-1].translate(NUCL_COMP),
        record.quals[::-1])


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]


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


if __name__ == '__main__':

    parser = ArgumentParser(
        description='Reverse complement',
        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(
        '-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()

    for f in args.input:
        with MultiReader(f) as F:
            for record in F:
                rc(record).print(wrap=args.wrap)
        f.close()