import os.path
import platform
from pip._vendor.rich.containers import Lines
import textwrap
from abc import ABC, abstractmethod
from typing import Any, Dict, Iterable, List, Optional, Set, Tuple, Type, Union
from pip._vendor.pygments.lexer import Lexer
from pip._vendor.pygments.lexers import get_lexer_by_name, guess_lexer_for_filename
from pip._vendor.pygments.style import Style as PygmentsStyle
from pip._vendor.pygments.styles import get_style_by_name
from pip._vendor.pygments.token import (
Comment,
Error,
Generic,
Keyword,
Name,
Number,
Operator,
String,
Token,
Whitespace,
)
from pip._vendor.pygments.util import ClassNotFound
from ._loop import loop_first
from .color import Color, blend_rgb
from .console import Console, ConsoleOptions, JustifyMethod, RenderResult
from .jupyter import JupyterMixin
from .measure import Measurement
from .segment import Segment
from .style import Style
from .text import Text
TokenType = Tuple[str, ...]
WINDOWS = platform.system() == "Windows"
DEFAULT_THEME = "monokai"
# The following styles are based on https://github.com/pygments/pygments/blob/master/pygments/formatters/terminal.py
# A few modifications were made
ANSI_LIGHT: Dict[TokenType, Style] = {
Token: Style(),
Whitespace: Style(color="white"),
Comment: Style(dim=True),
Comment.Preproc: Style(color="cyan"),
Keyword: Style(color="blue"),
Keyword.Type: Style(color="cyan"),
Operator.Word: Style(color="magenta"),
Name.Builtin: Style(color="cyan"),
Name.Function: Style(color="green"),
Name.Namespace: Style(color="cyan", underline=True),
Name.Class: Style(color="green", underline=True),
Name.Exception: Style(color="cyan"),
Name.Decorator: Style(color="magenta", bold=True),
Name.Variable: Style(color="red"),
Name.Constant: Style(color="red"),
Name.Attribute: Style(color="cyan"),
Name.Tag: Style(color="bright_blue"),
String: Style(color="yellow"),
Number: Style(color="blue"),
Generic.Deleted: Style(color="bright_red"),
Generic.Inserted: Style(color="green"),
Generic.Heading: Style(bold=True),
Generic.Subheading: Style(color="magenta", bold=True),
Generic.Prompt: Style(bold=True),
Generic.Error: Style(color="bright_red"),
Error: Style(color="red", underline=True),
}
ANSI_DARK: Dict[TokenType, Style] = {
Token: Style(),
Whitespace: Style(color="bright_black"),
Comment: Style(dim=True),
Comment.Preproc: Style(color="bright_cyan"),
Keyword: Style(color="bright_blue"),
Keyword.Type: Style(color="bright_cyan"),
Operator.Word: Style(color="bright_magenta"),
Name.Builtin: Style(color="bright_cyan"),
Name.Function: Style(color="bright_green"),
Name.Namespace: Style(color="bright_cyan", underline=True),
Name.Class: Style(color="bright_green", underline=True),
Name.Exception: Style(color="bright_cyan"),
Name.Decorator: Style(color="bright_magenta", bold=True),
Name.Variable: Style(color="bright_red"),
Name.Constant: Style(color="bright_red"),
Name.Attribute: Style(color="bright_cyan"),
Name.Tag: Style(color="bright_blue"),
String: Style(color="yellow"),
Number: Style(color="bright_blue"),
Generic.Deleted: Style(color="bright_red"),
Generic.Inserted: Style(color="bright_green"),
Generic.Heading: Style(bold=True),
Generic.Subheading: Style(color="bright_magenta", bold=True),
Generic.Prompt: Style(bold=True),
Generic.Error: Style(color="bright_red"),
Error: Style(color="red", underline=True),
}
RICH_SYNTAX_THEMES = {"ansi_light": ANSI_LIGHT, "ansi_dark": ANSI_DARK}
class SyntaxTheme(ABC):
"""Base class for a syntax theme."""
@abstractmethod
def get_style_for_token(self, token_type: TokenType) -> Style:
"""Get a style for a given Pygments token."""
raise NotImplementedError # pragma: no cover
@abstractmethod
def get_background_style(self) -> Style:
"""Get the background color."""
raise NotImplementedError # pragma: no cover
class PygmentsSyntaxTheme(SyntaxTheme):
"""Syntax theme that delegates to Pygments theme."""
def __init__(self, theme: Union[str, Type[PygmentsStyle]]) -> None:
self._style_cache: Dict[TokenType, Style] = {}
if isinstance(theme, str):
try:
self._pygments_style_class = get_style_by_name(theme)
except ClassNotFound:
self._pygments_style_class = get_style_by_name("default")
else:
self._pygments_style_class = theme
self._background_color = self._pygments_style_class.background_color
self._background_style = Style(bgcolor=self._background_color)
def get_style_for_token(self, token_type: TokenType) -> Style:
"""Get a style from a Pygments class."""
try:
return self._style_cache[token_type]
except KeyError:
try:
pygments_style = self._pygments_style_class.style_for_token(token_type)
except KeyError:
style = Style.null()
else:
color = pygments_style["color"]
bgcolor = pygments_style["bgcolor"]
style = Style(
color="#" + color if color else "#000000",
bgcolor="#" + bgcolor if bgcolor else self._background_color,
bold=pygments_style["bold"],
italic=pygments_style["italic"],
underline=pygments_style["underline"],
)
self._style_cache[token_type] = style
return style
def get_background_style(self) -> Style:
return self._background_style
class ANSISyntaxTheme(SyntaxTheme):
"""Syntax theme to use standard colors."""
def __init__(self, style_map: Dict[TokenType, Style]) -> None:
self.style_map = style_map
self._missing_style = Style.null()
self._background_style = Style.null()
self._style_cache: Dict[TokenType, Style] = {}
def get_style_for_token(self, token_type: TokenType) -> Style:
"""Look up style in the style map."""
try:
return self._style_cache[token_type]
except KeyError:
# Styles form a hierarchy
# We need to go from most to least specific
# e.g. ("foo", "bar", "baz") to ("foo", "bar") to ("foo",)
get_style = self.style_map.get
token = tuple(token_type)
style = self._missing_style
while token:
_style = get_style(token)
if _style is not None:
style = _style
break
token = token[:-1]
self._style_cache[token_type] = style
return style
def get_background_style(self) -> Style:
return self._background_style
class Syntax(JupyterMixin):
"""Construct a Syntax object to render syntax highlighted code.
Args:
code (str): Code to highlight.
lexer (Lexer | str): Lexer to use (see https://pygments.org/docs/lexers/)
theme (str, optional): Color theme, aka Pygments style (see https://pygments.org/docs/styles/#getting-a-list-of-available-styles). Defaults to "monokai".
dedent (bool, optional): Enable stripping of initial whitespace. Defaults to False.
line_numbers (bool, optional): Enable rendering of line numbers. Defaults to False.
start_line (int, optional): Starting number for line numbers. Defaults to 1.
line_range (Tuple[int, int], optional): If given should be a tuple of the start and end line to render.
highlight_lines (Set[int]): A set of line numbers to highlight.
code_width: Width of code to render (not including line numbers), or ``None`` to use all available width.
tab_size (int, optional): Size of tabs. Defaults to 4.
word_wrap (bool, optional): Enable word wrapping.
background_color (str, optional): Optional background color, or None to use theme color. Defaults to None.
indent_guides (bool, optional): Show indent guides. Defaults to False.
"""
_pygments_style_class: Type[PygmentsStyle]
_theme: SyntaxTheme
@classmethod
def get_theme(cls, name: Union[str, SyntaxTheme]) -> SyntaxTheme:
"""Get a syntax theme instance."""
if isinstance(name, SyntaxTheme):
return name
theme: SyntaxTheme
if name in RICH_SYNTAX_THEMES:
theme = ANSISyntaxTheme(RICH_SYNTAX_THEMES[name])
else:
theme = PygmentsSyntaxTheme(name)
return theme
def __init__(
self,
code: str,
lexer: Union[Lexer, str],
*,
theme: Union[str, SyntaxTheme] = DEFAULT_THEME,
dedent: bool = False,
line_numbers: bool = False,
start_line: int = 1,
line_range: Optional[Tuple[int, int]] = None,
highlight_lines: Optional[Set[int]] = None,
code_width: Optional[int] = None,
tab_size: int = 4,
word_wrap: bool = False,
background_color: Optional[str] = None,
indent_guides: bool = False,
) -> None:
self.code = code
self._lexer = lexer
self.dedent = dedent
self.line_numbers = line_numbers
self.start_line = start_line
self.line_range = line_range
self.highlight_lines = highlight_lines or set()
self.code_width = code_width
self.tab_size = tab_size
self.word_wrap = word_wrap
self.background_color = background_color
self.background_style = (
Style(bgcolor=background_color) if background_color else Style()
)
self.indent_guides = indent_guides
self._theme = self.get_theme(theme)
@classmethod
def from_path(
cls,
path: str,
encoding: str = "utf-8",
theme: Union[str, SyntaxTheme] = DEFAULT_THEME,
dedent: bool = False,
line_numbers: bool = False,
line_range: Optional[Tuple[int, int]] = None,
start_line: int = 1,
highlight_lines: Optional[Set[int]] = None,
code_width: Optional[int] = None,
tab_size: int = 4,
word_wrap: bool = False,
background_color: Optional[str] = None,
indent_guides: bool = False,
) -> "Syntax":
"""Construct a Syntax object from a file.
Args:
path (str): Path to file to highlight.
encoding (str): Encoding of file.
theme (str, optional): Color theme, aka Pygments style (see https://pygments.org/docs/styles/#getting-a-list-of-available-styles). Defaults to "emacs".
dedent (bool, optional): Enable stripping of initial whitespace. Defaults to True.
line_numbers (bool, optional): Enable rendering of line numbers. Defaults to False.
start_line (int, optional): Starting number for line numbers. Defaults to 1.
line_range (Tuple[int, int], optional): If given should be a tuple of the start and end line to render.
highlight_lines (Set[int]): A set of line numbers to highlight.
code_width: Width of code to render (not including line numbers), or ``None`` to use all available width.
tab_size (int, optional): Size of tabs. Defaults to 4.
word_wrap (bool, optional): Enable word wrapping of code.
background_color (str, optional): Optional background color, or None to use theme color. Defaults to None.
indent_guides (bool, optional): Show indent guides. Defaults to False.
Returns:
[Syntax]: A Syntax object that may be printed to the console
"""
with open(path, "rt", encoding=encoding) as code_file:
code = code_file.read()
lexer = None
lexer_name = "default"
try:
_, ext = os.path.splitext(path)
if ext:
extension = ext.lstrip(".").lower()
lexer = get_lexer_by_name(extension)
lexer_name = lexer.name
except ClassNotFound:
pass
if lexer is None:
try:
lexer_name = guess_lexer_for_filename(path, code).name
except ClassNotFound:
pass
return cls(
code,
lexer_name,
theme=theme,
dedent=dedent,
line_numbers=line_numbers,
line_range=line_range,
start_line=start_line,
highlight_lines=highlight_lines,
code_width=code_width,
tab_size=tab_size,
word_wrap=word_wrap,
background_color=background_color,
indent_guides=indent_guides,
)
def _get_base_style(self) -> Style:
"""Get the base style."""
default_style = self._theme.get_background_style() + self.background_style
return default_style
def _get_token_color(self, token_type: TokenType) -> Optional[Color]:
"""Get a color (if any) for the given token.
Args:
token_type (TokenType): A token type tuple from Pygments.
Returns:
Optional[Color]: Color from theme, or None for no color.
"""
style = self._theme.get_style_for_token(token_type)
return style.color
@property
def lexer(self) -> Optional[Lexer]:
"""The lexer for this syntax, or None if no lexer was found.
Tries to find the lexer by name if a string was passed to the constructor.
"""
if isinstance(self._lexer, Lexer):
return self._lexer
try:
return get_lexer_by_name(
self._lexer,
stripnl=False,
ensurenl=True,
tabsize=self.tab_size,
)
except ClassNotFound:
return None
def highlight(
self, code: str, line_range: Optional[Tuple[int, int]] = None
) -> Text:
"""Highlight code and return a Text instance.
Args:
code (str): Code to highlight.
line_range(Tuple[int, int], optional): Optional line range to highlight.
Returns:
Text: A text instance containing highlighted syntax.
"""
base_style = self._get_base_style()
justify: JustifyMethod = (
"default" if base_style.transparent_background else "left"
)
text = Text(
justify=justify,
style=base_style,
tab_size=self.tab_size,
no_wrap=not self.word_wrap,
)
_get_theme_style = self._theme.get_style_for_token
lexer = self.lexer
if lexer is None:
text.append(code)
else:
if line_range:
# More complicated path to only stylize a portion of the code
# This speeds up further operations as there are less spans to process
line_start, line_end = line_range
def line_tokenize() -> Iterable[Tuple[Any, str]]:
"""Split tokens to one per line."""
assert lexer
for token_type, token in lexer.get_tokens(code):
while token:
line_token, new_line, token = token.partition("\n")
yield token_type, line_token + new_line
def tokens_to_spans() -> Iterable[Tuple[str, Optional[Style]]]:
"""Convert tokens to spans."""
tokens = iter(line_tokenize())
line_no = 0
_line_start = line_start - 1
# Skip over tokens until line start
while line_no < _line_start:
_token_type, token = next(tokens)
yield (token, None)
if token.endswith("\n"):
line_no += 1
# Generate spans until line end
for token_type, token in tokens:
yield (token, _get_theme_style(token_type))
if token.endswith("\n"):
line_no += 1
if line_no >= line_end:
break
text.append_tokens(tokens_to_spans())
else:
text.append_tokens(
(token, _get_theme_style(token_type))
for token_type, token in lexer.get_tokens(code)
)
if self.background_color is not None:
text.stylize(f"on {self.background_color}")
return text
def _get_line_numbers_color(self, blend: float = 0.3) -> Color:
background_style = self._theme.get_background_style() + self.background_style
background_color = background_style.bgcolor
if background_color is None or background_color.is_system_defined:
return Color.default()
foreground_color = self._get_token_color(Token.Text)
if foreground_color is None or foreground_color.is_system_defined:
return foreground_color or Color.default()
new_color = blend_rgb(
background_color.get_truecolor(),
foreground_color.get_truecolor(),
cross_fade=blend,
)
return Color.from_triplet(new_color)
@property
def _numbers_column_width(self) -> int:
"""Get the number of characters used to render the numbers column."""
column_width = 0
if self.line_numbers:
column_width = len(str(self.start_line + self.code.count("\n"))) + 2
return column_width
def _get_number_styles(self, console: Console) -> Tuple[Style, Style, Style]:
"""Get background, number, and highlight styles for line numbers."""
background_style = self._get_base_style()
if background_style.transparent_background:
return Style.null(), Style(dim=True), Style.null()
if console.color_system in ("256", "truecolor"):
number_style = Style.chain(
background_style,
self._theme.get_style_for_token(Token.Text),
Style(color=self._get_line_numbers_color()),
self.background_style,
)
highlight_number_style = Style.chain(
background_style,
self._theme.get_style_for_token(Token.Text),
Style(bold=True, color=self._get_line_numbers_color(0.9)),
self.background_style,
)
else:
number_style = background_style + Style(dim=True)
highlight_number_style = background_style + Style(dim=False)
return background_style, number_style, highlight_number_style
def __rich_measure__(
self, console: "Console", options: "ConsoleOptions"
) -> "Measurement":
if self.code_width is not None:
width = self.code_width + self._numbers_column_width
return Measurement(self._numbers_column_width, width)
return Measurement(self._numbers_column_width, options.max_width)
def __rich_console__(
self, console: Console, options: ConsoleOptions
) -> RenderResult:
transparent_background = self._get_base_style().transparent_background
code_width = (
(
(options.max_width - self._numbers_column_width - 1)
if self.line_numbers
else options.max_width
)
if self.code_width is None
else self.code_width
)
line_offset = 0
if self.line_range:
start_line, end_line = self.line_range
line_offset = max(0, start_line - 1)
ends_on_nl = self.code.endswith("\n")
code = self.code if ends_on_nl else self.code + "\n"
code = textwrap.dedent(code) if self.dedent else code
code = code.expandtabs(self.tab_size)
text = self.highlight(code, self.line_range)
(
background_style,
number_style,
highlight_number_style,
) = self._get_number_styles(console)
if not self.line_numbers and not self.word_wrap and not self.line_range:
if not ends_on_nl:
text.remove_suffix("\n")
# Simple case of just rendering text
style = (
self._get_base_style()
+ self._theme.get_style_for_token(Comment)
+ Style(dim=True)
+ self.background_style
)
if self.indent_guides and not options.ascii_only:
text = text.with_indent_guides(self.tab_size, style=style)
text.overflow = "crop"
if style.transparent_background:
yield from console.render(
text, options=options.update(width=code_width)
)
else:
syntax_lines = console.render_lines(
text,
options.update(width=code_width, height=None),
style=self.background_style,
pad=True,
new_lines=True,
)
for syntax_line in syntax_lines:
yield from syntax_line
return
lines: Union[List[Text], Lines] = text.split("\n", allow_blank=ends_on_nl)
if self.line_range:
lines = lines[line_offset:end_line]
if self.indent_guides and not options.ascii_only:
style = (
self._get_base_style()
+ self._theme.get_style_for_token(Comment)
+ Style(dim=True)
+ self.background_style
)
lines = (
Text("\n")
.join(lines)
.with_indent_guides(self.tab_size, style=style)
.split("\n", allow_blank=True)
)
numbers_column_width = self._numbers_column_width
render_options = options.update(width=code_width)
highlight_line = self.highlight_lines.__contains__
_Segment = Segment
padding = _Segment(" " * numbers_column_width + " ", background_style)
new_line = _Segment("\n")
line_pointer = "> " if options.legacy_windows else "❱ "
for line_no, line in enumerate(lines, self.start_line + line_offset):
if self.word_wrap:
wrapped_lines = console.render_lines(
line,
render_options.update(height=None),
style=background_style,
pad=not transparent_background,
)
else:
segments = list(line.render(console, end=""))
if options.no_wrap:
wrapped_lines = [segments]
else:
wrapped_lines = [
_Segment.adjust_line_length(
segments,
render_options.max_width,
style=background_style,
pad=not transparent_background,
)
]
if self.line_numbers:
for first, wrapped_line in loop_first(wrapped_lines):
if first:
line_column = str(line_no).rjust(numbers_column_width - 2) + " "
if highlight_line(line_no):
yield _Segment(line_pointer, Style(color="red"))
yield _Segment(line_column, highlight_number_style)
else:
yield _Segment(" ", highlight_number_style)
yield _Segment(line_column, number_style)
else:
yield padding
yield from wrapped_line
yield new_line
else:
for wrapped_line in wrapped_lines:
yield from wrapped_line
yield new_line
if __name__ == "__main__": # pragma: no cover
import argparse
import sys
parser = argparse.ArgumentParser(
description="Render syntax to the console with Rich"
)
parser.add_argument(
"path",
metavar="PATH",
help="path to file, or - for stdin",
)
parser.add_argument(
"-c",
"--force-color",
dest="force_color",
action="store_true",
default=None,
help="force color for non-terminals",
)
parser.add_argument(
"-i",
"--indent-guides",
dest="indent_guides",
action="store_true",
default=False,
help="display indent guides",
)
parser.add_argument(
"-l",
"--line-numbers",
dest="line_numbers",
action="store_true",
help="render line numbers",
)
parser.add_argument(
"-w",
"--width",
type=int,
dest="width",
default=None,
help="width of output (default will auto-detect)",
)
parser.add_argument(
"-r",
"--wrap",
dest="word_wrap",
action="store_true",
default=False,
help="word wrap long lines",
)
parser.add_argument(
"-s",
"--soft-wrap",
action="store_true",
dest="soft_wrap",
default=False,
help="enable soft wrapping mode",
)
parser.add_argument(
"-t", "--theme", dest="theme", default="monokai", help="pygments theme"
)
parser.add_argument(
"-b",
"--background-color",
dest="background_color",
default=None,
help="Override background color",
)
parser.add_argument(
"-x",
"--lexer",
default="default",
dest="lexer_name",
help="Lexer name",
)
args = parser.parse_args()
from pip._vendor.rich.console import Console
console = Console(force_terminal=args.force_color, width=args.width)
if args.path == "-":
code = sys.stdin.read()
syntax = Syntax(
code=code,
lexer=args.lexer_name,
line_numbers=args.line_numbers,
word_wrap=args.word_wrap,
theme=args.theme,
background_color=args.background_color,
indent_guides=args.indent_guides,
)
else:
syntax = Syntax.from_path(
args.path,
line_numbers=args.line_numbers,
word_wrap=args.word_wrap,
theme=args.theme,
background_color=args.background_color,
indent_guides=args.indent_guides,
)
console.print(syntax, soft_wrap=args.soft_wrap)