Python excepthook with support for zipimport

exhook.py

import sys
from types import TracebackType
from termcolor import colored
import os
import zipfile
import math

LINE_LOOKAROUND = 3
SHOW_LOCALS_AT_TOP = True


def excepthook(exception_type: type[BaseException], exception: BaseException, tb: TracebackType | None):
    t = tb
    print("Traceback (most recent call last):")

    def fetch_lines(lines: list[str], line: int):
        line_slice = slice(max(0, line-LINE_LOOKAROUND),
                           min(len(lines)-1, line+LINE_LOOKAROUND)+1)
        return map(lambda x: x.rstrip(), lines[line_slice]), line_slice

    def fetch_codes(filename: str, line: int):
        if os.path.exists(filename):
            with open(filename) as f:
                lines = f.readlines()
                return fetch_lines(lines, line)
        zip_file = os.path.dirname(filename)
        while len(zip_file) > 0 and not zipfile.is_zipfile(zip_file):
            zip_file = os.path.dirname(zip_file)
        if len(zip_file) == 0:
            return None, None
        file_in_zip = filename[len(zip_file)+1:]
        with zipfile.ZipFile(zip_file) as zf:
            return fetch_lines(zf.read(file_in_zip).decode().splitlines(), line)

    top_tb = True
    while t is not None:
        frame = t.tb_frame
        filename = os.path.abspath(frame.f_code.co_filename)
        dir = os.path.dirname(filename)
        fname = os.path.basename(filename)
        print("  {dir}/{fname}:{line} in {module}".format(
            dir=colored(dir, color='yellow'),
            fname=colored(fname, color='light_yellow'),
            line=colored(t.tb_lineno, color='blue'),
            module=colored(frame.f_code.co_qualname, color='green'),
        ))
        lines, line_slice = fetch_codes(filename, t.tb_lineno)
        if lines is not None:
            line_nos = list(range(line_slice.start, line_slice.stop))
            padding = len(str(max(line_nos)))
            for (line_no, line) in zip(line_nos, lines):
                is_current_line = line_no+1 == t.tb_lineno
                print('    {indicator} {line_no:<{padding}}: {line}'.format(
                    indicator=colored(
                        '>>', color='red') if is_current_line else '  ',
                    line_no=colored(
                        line_no+1, color=None if not is_current_line else 'red'),
                    padding=padding,
                    line=line
                ))

        if SHOW_LOCALS_AT_TOP and top_tb:
            keys = list(filter(lambda key: not key.startswith(
                '__') and not key.endswith('__'), t.tb_frame.f_locals.keys()))
            values = list(map(lambda key: str(
                t.tb_frame.f_locals[key]), keys))
            if len(keys) > 0:
                max_name_len = max(map(lambda x: len(x), keys))
                max_value_len = max(map(lambda x: len(x), values))
                separator_chars = max_name_len + max_value_len + 3 + 4 # ' = ' and margin
                separator_chars = math.ceil(separator_chars/2) * 2
                title = ' locals '
                print('      {sp}{title}{sp}'.format(
                    sp='-' * ((separator_chars - len(title))//2),
                    title=title,))
                for name, value in zip(keys, values):
                    print(f'        {name:>{max_name_len}} = {value}')
                print('      ' + ('-'*separator_chars))

        t = t.tb_next
        top_tb = False

    print()
    print('{ty}: {info}'.format(
        ty=colored(exception_type.__name__, color='red'),
        info=exception
    ))


def install():
    sys.excepthook = excepthook