Python highlighter

#!/usr/bin/env python3
"""Render Python source as safe, dependency-free syntax-highlighted HTML.

The public :func:`render_python_source` function returns an HTML fragment that
can be placed inside a ``<code>`` element.  Python's standard-library
``tokenize`` module supplies the token boundaries, while all source text is
escaped before it is emitted.
"""

import builtins
import html
import io
import keyword
import re
import token
import tokenize
from typing import List, Optional, Sequence, Tuple


PYTHON_SOURCE_CSS = """
.python-source {
  color: var(--syntax-name, var(--fn-syntax-name, #243247));
  tab-size: 4;
}
.python-source .syntax-keyword {
  color: var(--syntax-keyword, var(--fn-syntax-keyword, #8f3f71));
  font-weight: 650;
}
.python-source .syntax-builtin {
  color: var(--syntax-builtin, var(--fn-syntax-builtin, #315f8a));
}
.python-source .syntax-name {
  color: var(--syntax-name, var(--fn-syntax-name, #243247));
}
.python-source .syntax-string {
  color: var(--syntax-string, var(--fn-syntax-string, #386747));
}
.python-source .syntax-number {
  color: var(--syntax-number, var(--fn-syntax-number, #8c511d));
}
.python-source .syntax-comment {
  color: var(--syntax-comment, var(--fn-syntax-comment, #627063));
  font-style: italic;
}
.python-source .syntax-operator {
  color: var(--syntax-operator, var(--fn-syntax-operator, #6f4b7d));
}
.python-source .syntax-decorator {
  color: var(--syntax-decorator, var(--fn-syntax-decorator, #86561e));
  font-weight: 650;
}
.python-source .source-line {
  display: block;
  min-height: var(--fn-code-leading, 1em);
  position: relative;
}
.python-source.has-line-numbers .source-line {
  padding-left: var(--fn-code-gutter, 3.75rem);
}
.python-source .source-line-number {
  color: var(--fn-line-number, #6d756f);
  font-size: 0.82em;
  font-variant-numeric: tabular-nums;
  left: 0;
  position: absolute;
  text-align: right;
  text-decoration: none;
  width: var(--fn-line-number-width, 2.7rem);
}
.python-source .source-line-number::before { content: attr(data-line-number); }
.python-source .source-line-number:hover { color: currentColor; text-decoration: underline; }
""".strip()


_BUILTIN_NAMES = frozenset(dir(builtins))
_SAFE_PREFIX = re.compile(r"[^A-Za-z0-9_-]+")
_PHYSICAL_LINE = re.compile(r"([^\r\n]*)(\r\n|\r|\n|$)")
_TokenRange = Tuple[int, int, str]


def _line_starts(source: str) -> List[int]:
    starts = [0]
    starts.extend(index + 1 for index, character in enumerate(source) if character == "\n")
    return starts


def _offset(starts: Sequence[int], position: Tuple[int, int], source_length: int) -> int:
    row, column = position
    line_index = row - 1
    if line_index < 0 or line_index >= len(starts):
        return source_length
    return min(starts[line_index] + column, source_length)


def _ordinary_class(token_type: int, value: str) -> Optional[str]:
    if token_type == token.NAME:
        if keyword.iskeyword(value):
            return "syntax-keyword"
        if value in _BUILTIN_NAMES:
            return "syntax-builtin"
        return "syntax-name"
    if token_type == token.STRING:
        return "syntax-string"
    if token_type == token.NUMBER:
        return "syntax-number"
    if token_type == tokenize.COMMENT:
        return "syntax-comment"
    if token_type == token.OP:
        return "syntax-operator"
    return None


def _token_ranges(source: str) -> List[_TokenRange]:
    starts = _line_starts(source)
    ranges: List[_TokenRange] = []
    at_statement_start = True
    in_decorator_reference = False

    for item in tokenize.generate_tokens(io.StringIO(source).readline):
        token_type, value, start, end, _ = item
        css_class: Optional[str]

        if token_type in (token.INDENT, token.DEDENT):
            at_statement_start = True
            continue
        if token_type in (token.NEWLINE, tokenize.NL):
            if token_type == token.NEWLINE:
                at_statement_start = True
            in_decorator_reference = False
            continue

        is_decorator_at = token_type == token.OP and value == "@" and at_statement_start
        if is_decorator_at:
            css_class = "syntax-decorator"
            in_decorator_reference = True
        elif in_decorator_reference and (
            token_type == token.NAME or (token_type == token.OP and value == ".")
        ):
            css_class = "syntax-decorator"
        else:
            if in_decorator_reference:
                in_decorator_reference = False
            css_class = _ordinary_class(token_type, value)

        if css_class:
            range_start = _offset(starts, start, len(source))
            range_end = _offset(starts, end, len(source))
            if range_end > range_start:
                ranges.append((range_start, range_end, css_class))

        if token_type not in (
            token.ENCODING,
            token.ENDMARKER,
            token.INDENT,
            token.DEDENT,
            tokenize.COMMENT,
        ):
            at_statement_start = False

    return ranges


def _render_interval(
    source: str,
    start: int,
    end: int,
    ranges: Sequence[_TokenRange],
    range_index: int = 0,
) -> str:
    rendered: List[str] = []
    cursor = start
    for index in range(range_index, len(ranges)):
        range_start, range_end, css_class = ranges[index]
        if range_end <= start:
            continue
        if range_start >= end:
            break
        token_start = max(range_start, start)
        token_end = min(range_end, end)
        if token_start > cursor:
            rendered.append(html.escape(source[cursor:token_start], quote=True))
        if token_end > token_start:
            rendered.append(
                f'<span class="{css_class}">'
                f"{html.escape(source[token_start:token_end], quote=True)}</span>"
            )
            cursor = token_end
    if cursor < end:
        rendered.append(html.escape(source[cursor:end], quote=True))
    return "".join(rendered)


def _render_reader_lines(source: str, ranges: Sequence[_TokenRange]) -> str:
    rendered: List[str] = []
    range_index = 0

    for line_number, match in enumerate(_PHYSICAL_LINE.finditer(source), start=1):
        if match.start() == len(source):
            break
        line_start, line_end = match.span(1)
        while range_index < len(ranges) and ranges[range_index][1] <= line_start:
            range_index += 1
        rendered.append(
            f'<span class="code-line" id="L{line_number}" data-line="{line_number}">'
            f"{_render_interval(source, line_start, line_end, ranges, range_index)}"
            f"</span>{match.group(2)}"
        )

    return "".join(rendered)


def _safe_prefix(value: str) -> str:
    prefix = _SAFE_PREFIX.sub("-", value).strip("-")
    return prefix or "python-source"


def _render_numbered(source: str, ranges: Sequence[_TokenRange], anchor_prefix: str) -> str:
    prefix = _safe_prefix(anchor_prefix)
    rendered: List[str] = []
    line_start = 0
    line_number = 1

    for newline_index in (
        [index for index, character in enumerate(source) if character == "\n"] + [len(source)]
    ):
        line_end = newline_index
        newline = ""
        if newline_index < len(source):
            newline = "\n"
            if line_end > line_start and source[line_end - 1] == "\r":
                line_end -= 1
                newline = "\r\n"
        if line_start == len(source) and not newline:
            break
        identifier = f"{prefix}-L{line_number}"
        rendered.append(
            f'<span class="source-line" id="{identifier}">'
            f'<a class="source-line-number" href="#{identifier}" '
            f'data-line-number="{line_number}" aria-label="Line {line_number}" '
            'tabindex="-1"></a>'
            f"{_render_interval(source, line_start, line_end, ranges)}</span>{newline}"
        )
        line_start = newline_index + 1
        line_number += 1

    return "".join(rendered)


def render_python_source(
    source: str,
    *,
    line_numbers: bool = False,
    anchor_prefix: str = "python-source",
    compact: bool = False,
    reader_lines: bool = False,
) -> str:
    """Return a safe syntax-highlighted HTML fragment for *source*.

    Whitespace and newlines from the input are emitted exactly.  When
    ``line_numbers`` is true, every physical source line receives a stable,
    accessible anchor.  ``reader_lines`` instead wraps each physical line in a
    compact span with its own ``L1``-style anchor and line-number attribute for
    a CSS gutter.  This mode takes precedence over ``line_numbers``; it emits no
    anchor elements, and keeps literal line endings outside the wrappers.
    Compact mode leaves ordinary identifiers and operators
    unwrapped because the code element already supplies their base color.  This
    keeps large public source pages small without changing their text. Malformed
    or incomplete Python falls back to fully escaped, unhighlighted source
    instead of failing the surrounding render.
    """

    try:
        ranges = _token_ranges(source)
    except (tokenize.TokenError, IndentationError, SyntaxError):
        ranges = []

    if compact:
        ranges = [
            value
            for value in ranges
            if value[2] not in {"syntax-name", "syntax-operator"}
        ]

    if reader_lines:
        return _render_reader_lines(source, ranges)
    if line_numbers:
        return _render_numbered(source, ranges, anchor_prefix)
    return _render_interval(source, 0, len(source), ranges)