head_tail_count.py

#!/usr/bin/env python3

import sys
from contextlib import redirect_stdout
from pathlib import Path
from collections import deque


def help():
    with redirect_stdout(sys.stderr):
        print('This programm counts the words in a file')
        print('The command head counts the words in the first n lines')
        print('The command tail counts the words in the last n lines\n')
        print(sys.argv[0], '<head|tail> <file> <n-lines>')
    sys.exit(1)


def get_args():
    try:
        exe, cmd, file, n = sys.argv
    except ValueError:
        help()
    if not cmd in ('head', 'tail'):
        help()
    try:
        n = int(n)
    except ValueError:
        help()
    try:
        file = Path(file)
    except ValueError:
        print('The path is invalid', file=sys.stderr)
        sys.exit(2)
    if not file.exists():
        print('The file does not exist', file=sys.stderr)
        sys.exit(3)
    cmd = globals().get(cmd)
    return cmd, file, n


def head(iterable, n=1):
    for element, _ in zip(iterable, range(n)):
        yield element


def tail(iterable, n=1):
    hist = deque(maxlen=n)
    for element in iterable:
        hist.append(element)
    yield from hist


def count_words(func, file, n):
    word_count = 0
    with open(file) as fd:
        for line in func(fd, n):
            word_count += len(line.split())
    print(f'{word_count} words in {n} line(s)')


if __name__ == '__main__':
    cmd, file, n = get_args()
    count_words(cmd, file, n)