grep2.py

#!/usr/bin/env python3
from typing import List
import sys
import re
import argparse


def parse_args(args_str: List[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument('needle', type=str)
    parser.add_argument('files', nargs='*')
    parser.add_argument('-E', dest='regex', action='store_true')
    parser.add_argument('-c', dest='counting', action='store_true')
    args = parser.parse_args(args_str)
    return args


def print_in_files(files_amount: bool, file: str, parametr: str) -> None:
    print(f'{file}:{parametr}' if files_amount else parametr)


def output_in_files(files_amount: bool, lines: List[str], file: str) -> None:
    for line in lines:
        print_in_files(files_amount, file, line)


def output_in_stdin(lines: List[str]) -> None:
    for line in lines:
        print(line)


def add_to_list(lines: List[str], line: str, needle: str) -> None:
    line = line.rstrip('\n')
    if re.search(needle, line):
        lines.append(line)


def making_list_in_file(file: str, needle: str) -> list:
    lines = []
    with open(file, 'r') as in_file:
        for line in in_file:
            add_to_list(lines, line, needle)
    return lines


def making_list_in_stdin(needle: str) -> list:
    lines = []
    for line in sys.stdin.readlines():
        add_to_list(lines, line, needle)
    return lines


def search_in_file(files: List[str], needle: str, counting: bool) -> None:
    for file in files:
        lines = making_list_in_file(file, needle)
        lines = [str(len(lines))] if counting else lines
        output_in_files(bool(len(files) > 1), lines, file)


def search_in_stdin(needle: str, counting: bool) -> None:
    lines = making_list_in_stdin(needle)
    lines = [str(len(lines))] if counting else lines
    output_in_stdin(lines)


def main(args_str: List[str]) -> None:
    args = parse_args(args_str)
    needle = args.needle if args.regex else re.escape(args.needle)
    if args.files:
        search_in_file(args.files, needle, args.counting)
    else:
        search_in_stdin(needle, args.counting)


if __name__ == '__main__':
    args = main(sys.argv[1:])