Receive regex expressions, reads from stdin, and split the output into different columns - each for every regex

split_grep.py

#!/usr/bin/env python

from typing import List
import typer  # pip install typer==0.6.1
import termios
import fcntl
import os
import struct
import sys
import re


RED_START = '\x1b[31m'
RED_END = '\x1b[0m'
SPLITTER = "|"


def get_win_size():
    '''
    Returns the terminal windows size.
    Taken from sbz gist: https://gist.github.com/sbz/bca2d1327910e0d8cadff1835625cbf2
    '''
    with open(os.ctermid(), 'r') as fd:
        packed = fcntl.ioctl(fd, termios.TIOCGWINSZ, struct.pack('HHHH', 0, 0, 0, 0))
        rows, cols, h_pixels, v_pixels = struct.unpack('HHHH', packed)
        return rows, cols, h_pixels, v_pixels


def split_color_lines(re_match, line, cols):
    start = re_match.start()
    end = re_match.end()
    color_line = line[:start] + RED_START + line[start:end] + RED_END + line[end:]
    i = 0
    lines = []
    while i < len(line):
        start_addition = 0
        if end < i:
            start_addition += len(RED_END) + len(RED_START)
            start_color = ""
        elif start < i:
            start_addition += len(RED_START)
            start_color = RED_START
        else:
            start_color = ""

        end_addition = 0
        if end < i + cols:
            end_addition += len(RED_END) + len(RED_START)
            end_color = ""
        elif start < i + cols:
            end_addition += len(RED_START)
            end_color = RED_END
        else:
            end_color = ""

        suffix = " "*(i + cols - len(line))
        lines.append(start_color + color_line[start_addition + i : end_addition + i + cols] + end_color + suffix)
        i += cols
    return lines

def handle_line(line: str, patterns: List[re.Pattern], col_size: int):
    # Remove the new line from the end
    line = line[:-1]
    prefix = ""
    for index, pattern in enumerate(patterns):
        re_match = pattern.search(line)
        if re_match:
            splited_lines = split_color_lines(re_match, line, col_size)
            cols_left = len(patterns) - index - 1
            suffix = cols_left * (SPLITTER + " "*col_size)
            for splited_line in splited_lines:
                print(prefix + splited_line + suffix)
            break
        else:
            prefix += " "*col_size + SPLITTER


def main(regex_list: List[str]):
    patterns: List[re.Pattern] = []
    for r in regex_list:
        patterns.append(re.compile(r))
    _, cols, _, _ = get_win_size()
    col_size = int(cols/len(patterns) - 1)

    while True:
        line = sys.stdin.readline()
        if not line:
            break

        handle_line(line, patterns, col_size)


if __name__ == "__main__":
    typer.run(main)