mx-fa-sortseqlen

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# vim:fileencoding=utf-8

"""
SORT FASTA BY SEQUENCE LENGTH
James B. Pease
"""
import os
import sys
import argparse
from pylib.mixcore import _LICENSE
from pylib.mixcore import fasta_iter


def generate_argparser():
    parser = argparse.ArgumentParser(
        description=__doc__,
        formatter_class=argparse.RawTextHelpFormatter,
        epilog=_LICENSE)
    parser.add_argument('fasta', type=os.path.abspath,
                        help="input fasta file")
    parser.add_argument('--out', type=argparse.FileType('w'),
                        default=sys.stdout, help="output fasta file")
    parser.add_argument('--ignore-chars', help=(
                        "Ignore the characters in the string when counting."
                        "For example, use '-' to ignore gaps or '-N' for gaps and N's."))
    parser.add_argument('--sort-labels-next', action="store_true", help=(
                        "After sorting by length, next sort by label."))
    return parser

def main(arguments=None):
    arguments = arguments if arguments is not None else sys.argv[1:]
    parser = generate_argparser()
    args = parser.parse_args(args=arguments)
    seqs = []
    for hdr, seq in fasta_iter(args.fasta):
        if args.ignore_chars:
            seqlen = sum((0 if x in args.ignore_chars else 1 for x in seq))
        else:
            seqlen = len(seq)
        seqs.append((seqlen, hdr, seq))
    if args.sort_labels_next is True:
        for _, hdr, seq in sorted(seqs):
            print(">{}\n{}".format(hdr, seq), file=args.out)
    else:
        for _, hdr, seq in sorted(seqs, key=lambda x:x[0]):
            print(">{}\n{}".format(hdr, seq), file=args.out)
    return ''

if __name__ == '__main__':
    main()