rio/rio/cli/nice_traceback.py

"""
Pretty-strings a traceback. The result looks very similar to Python's default,
but is colored and just tweaked in general.
"""

import dataclasses
import html
import io
import linecache
import sys
import traceback
import typing as t
from pathlib import Path

import revel


@dataclasses.dataclass
class FormatStyle:
    bold: str
    nobold: str
    dim: str
    nodim: str
    yellow: str
    noyellow: str
    red: str
    nored: str
    escape: t.Callable[[str], str]


def _handle_syntax_error(err: SyntaxError) -> traceback.FrameSummary:
    """
    Syntax errors are very special snowflakes and need separate treatment. This
    creates a FrameSummary for a SyntaxError, handling differences between
    Python versions.
    """
    filename = err.filename or "<unknown>"

    # TODO: Is there a better way to do this? Since Python obviously isn't
    # keeping the arguments consistent, this could potentially break every
    # single Python version.
    if sys.version_info < (3, 11):
        return traceback.FrameSummary(
            filename=filename,
            lineno=err.lineno,
            name="<module>",
            line=err.text,
            locals=None,
        )
    else:
        return traceback.FrameSummary(
            filename=filename,
            lineno=err.lineno,
            end_lineno=err.end_lineno,
            colno=err.offset,
            end_colno=err.end_offset,
            name="<module>",
            line=err.text,
            locals=None,
        )


def _format_single_exception_raw(
    out: t.IO[str],
    err: BaseException,
    *,
    include_header: bool,
    style: FormatStyle,
    relpath: Path | None,
    frame_filter: t.Callable[[traceback.FrameSummary], bool],
) -> None:
    """
    Format a single exception and write it to the output stream.
    """
    tb_list = traceback.extract_tb(err.__traceback__)

    # Syntax errors need special handling. Convert them to something that
    # behaves more like a regular one.
    if isinstance(err, SyntaxError):
        tb_list.append(_handle_syntax_error(err))

    # TODO: Add special handling for recursion errors. Instead of printing the
    # same frame 1000 times, print a message like "Last 5 frames repeated 200
    # times".

    # Lead-in
    if include_header:
        out.write(
            f"{style.dim}Traceback (most recent call last):{style.nodim}\n"
        )

    # Walk all frames and format them
    for frame in tb_list:
        assert frame.lineno is not None

        # Keep this frame?
        if not frame_filter(frame):
            continue

        # Make paths relative to `relpath` if they're contained within
        frame_path = Path(frame.filename)
        if relpath and frame_path.is_absolute():
            try:
                frame_path = frame_path.relative_to(relpath)
            except ValueError:
                pass

        # File location
        out.write(
            f"  {style.dim}File{style.nodim} {style.yellow}{style.escape(str(frame_path))}{style.noyellow}"
            f"{style.dim}, {style.nodim}line {frame.lineno}{style.noyellow}{style.dim}, in {style.escape(frame.name)}{style.nodim}\n"
        )

        # Display the source code from this line
        #
        # Insanely, the line in the frame has been stripped. Refetch it.
        source_line = linecache.getline(frame.filename, frame.lineno)
        if source_line:
            # If this is the last line, highlight the error
            if (
                frame is tb_list[-1]
                and hasattr(frame, "colno")
                and hasattr(frame, "end_colno")
                and frame.colno is not None  # type: ignore
                and frame.end_colno is not None  # type: ignore
            ):
                if (
                    hasattr(frame, "end_lineno")
                    and frame.end_lineno is not None  # type: ignore
                    and frame.end_lineno > frame.lineno  # type: ignore
                ):
                    end_col = len(source_line) - 1  # -1 to exclude the \n
                else:
                    end_col = frame.end_colno  # type: ignore

                before = style.escape(source_line[: frame.colno].lstrip())  # type: ignore
                error = style.escape(source_line[frame.colno : end_col])  # type: ignore
                after = style.escape(source_line[end_col:].rstrip())
                formatted_line = (
                    f"{before}{style.red}{error}{style.nored}{after}"
                )
            else:
                formatted_line = style.escape(source_line.strip())

            # Write the line
            out.write(f"    {formatted_line}\n")

    # Exception type and message
    out.write("\n")
    out.write(
        f"{style.bold}{style.red}{type(err).__name__}{style.nored}{style.nobold}"
    )

    error_message = str(err)

    if error_message:
        out.write(f"{style.bold}: {style.escape(error_message)}{style.nobold}")


def format_exception_raw(
    err: BaseException,
    *,
    style: FormatStyle,
    relpath: Path | None = None,
    frame_filter: t.Callable[[traceback.FrameSummary], bool] = lambda _: True,
) -> str:
    """
    Format an exception into a pretty string with the given style.
    """
    frame_filter = frame_filter or (lambda _: True)

    def format_inner(current_err: BaseException) -> None:
        # Chain to the cause or context if there is one
        if current_err.__cause__ is not None:
            format_inner(current_err.__cause__)
            out.write("\n\n")
            out.write(
                "The above exception was the direct cause of the following:\n\n"
            )
            include_header = False
        elif current_err.__context__ is not None:
            format_inner(current_err.__context__)
            out.write("\n\n")
            out.write(
                "During handling of the above exception, another exception occurred:\n\n"
            )
            include_header = False
        else:
            include_header = True

        # Format the exception
        _format_single_exception_raw(
            out,
            current_err,
            include_header=include_header,
            style=style,
            relpath=relpath,
            frame_filter=frame_filter,
        )

    # Format
    out = io.StringIO()
    format_inner(err)
    return out.getvalue()


def format_exception_revel(
    err: BaseException,
    *,
    relpath: Path | None = None,
    frame_filter: t.Callable[[traceback.FrameSummary], bool] = lambda _: True,
) -> str:
    """
    Format an exception using revel's styling.
    """

    # Prepare the style
    style = FormatStyle(
        bold="[bold]",
        nobold="[/]",
        dim="[dim]",
        nodim="[/]",
        yellow="[yellow]",
        noyellow="[/]",
        red="[red]",
        nored="[/]",
        escape=revel.escape,
    )

    # Format the exception
    return format_exception_raw(
        err,
        style=style,
        relpath=relpath,
        frame_filter=frame_filter,
    )


def format_exception_html(
    err: BaseException,
    *,
    relpath: Path | None = None,
    frame_filter: t.Callable[[traceback.FrameSummary], bool] = lambda _: True,
) -> str:
    """
    Format an exception into HTML with appropriate styling.
    """

    # Prepare the style
    style = FormatStyle(
        bold='<span class="rio-traceback-bold">',
        nobold="</span>",
        dim='<span class="rio-traceback-dim">',
        nodim="</span>",
        yellow='<span class="rio-traceback-yellow">',
        noyellow="</span>",
        red='<span class="rio-traceback-red">',
        nored="</span>",
        escape=html.escape,
    )

    # Format the exception
    result = format_exception_raw(
        err,
        style=style,
        relpath=relpath,
        frame_filter=frame_filter,
    )

    # HTML-ify newlines
    return result.replace("\n", "<br>")