lines

#!/usr/bin/env python3

"""
Print arbitrary lines from a file or list of files.
Default input stream is sys.stdin.

Usage:

   lines [-v] [-0] LINENO FILE1 [FILE2 FILE3 ...]
   cat FILE | mid LINENO

where LINENO is a comma-separated list of line selectors (indexed from 1). 
Valid selectors are:

- a single number (10)
- a range (1-10)
- an offset (10+9 to print lines 10-19)

If `-v` is present, it will print a UNIX `head`-style header above each
file, identifying its provenance. `-0` can be used to use zero indexing.

Example usage:

    seq 1 100 > t

    # pipeline
    $ cat t | lines 1,10,15-18,20+1
    1
    10
    15
    16
    17
    18
    20
    21

    # file list
    $ cat t | lines 1,10,15-18,20+1
    (same output)

    # zero offset
    $ cat t | lines -0 1,10,15,20+1
    2
    11
    16
    17
    18
    19
    21
    22

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

import os
import sys
import gzip
import argparse
from itertools import zip_longest


def main(args):
    groups = args.lines.split(",")
    wanted = set()

    def num(i):
        if args.zero:
            return int(i) + 1
        else:
            return int(i)

    for group in groups:
        try:
            if '-' in group:
                start, stop = map(num, group.split("-"))
                wanted = wanted.union(set(range(start, stop + 1)))
            elif '+' in group:
                start, stop = group.split("+")
                start = num(start)
                wanted = wanted.union(range(start, start + int(stop) + 1))
            else:
                wanted.add(num(group))
        except ValueError as e:
            print(f"Invalid group '{group}'", file=sys.stderr)
            sys.exit(1)

    toprint = sorted(wanted)

    for lineno, lines in enumerate(zip_longest(*args.files), 1):
        if len(toprint) == 0:
            break
        if lineno != toprint[0]:
            continue

        toprint.pop(0)

        for inputno, line in enumerate(lines):
            if line is not None:
                if args.verbose:
                    print('=== {}'.format(args.files[inputno].name), line.rstrip(), '', sep='\n')
                else:
                    print(line.rstrip(),)

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("Can't find file '{}'".format(filename))
        sys.exit(1)


if __name__ == '__main__':
    parser = argparse.ArgumentParser('Grabs a line from one or more files')
    parser.add_argument('--verbose', '-v', default=False, action='store_true',
			help='Display head-style file headers if there is more than one file')
    parser.add_argument("--zero", "-0", action="store_true",
                        help="Use 0-based indexing instead of 1-based")
    parser.add_argument('lines',
			help='The line number (N) or range (M-N or M:N) to display')
    parser.add_argument('files',
                        nargs='*',
                        type=smart_read,
                        default=[sys.stdin],
			help='The files to work with')
    args = parser.parse_args()

    main(args)