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_-]+")
_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]) -> str:
    rendered: List[str] = []
    cursor = start
    for range_start, range_end, css_class in ranges:
        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 _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,
) -> 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.  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 line_numbers:
        return _render_numbered(source, ranges, anchor_prefix)
    return _render_interval(source, 0, len(source), ranges)