zpaste

#!/usr/bin/env python3

"""
"paste" annoyingly only works with uncompressed files. This changes that.

Author: Matt Post <post@cs.jhu.edu
"""

import sys
import gzip
from itertools import zip_longest

def smart_read(filename: str):
    try:
        if filename == '-':
            return sys.stdin
        elif filename.endswith('.gz'):
            return gzip.open(filename, mode='rt', encoding='utf-8')
        else:
            return open(filename, mode='rt', encoding='utf-8')
    except FileNotFoundError as e:
        print(f"Fatal: no such file '{filename}'", file=sys.stderr)
        sys.exit(1)

def main(args):
    zip_func = zip if args.shortest else zip_longest
    for i, lines in enumerate(zip_func(*args.files), 1):
        if args.head != 0 and i > args.head:
            break
        lines = list(map(lambda x: x.rstrip(), lines))
        if args.line_numbers:
            lines.insert(0, str(i))
        print(args.delimiter.join(lines))

if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser('A better UNIX paste')
    parser.add_argument('-d', '--delimiter', default='\t', type=str, help='Delimiter to split on. Default: %(default)s.')
    parser.add_argument('--shortest', action='store_true', help='Stop when the first file runs out.')
    parser.add_argument('--head', type=int, default=0, help='Only output first HEAD lines. Default: %(default)s.')
    parser.add_argument('--line-numbers', '-l', action='store_true', help='Prepend line numbers as first column')
    parser.add_argument('files',
                        nargs='+',
                        type=smart_read,
                        default=[sys.stdin],
			help='The file streams to paste together')
    args = parser.parse_args()

    main(args)