import os.path import platform import textwrap from abc import ABC, abstractmethod from typing import Any, Dict, Iterable, Optional, Set, Tuple, Type, Union from pygments.lexers import get_lexer_by_name, guess_lexer_for_filename from pygments.style import Style as PygmentsStyle from pygments.styles import get_style_by_name from pygments.token import ( Comment, Error, Generic, Keyword, Name, Number, Operator, String, Token, Whitespace, ) from pygments.util import ClassNotFound from ._loop import loop_first from .color import Color, blend_rgb from .console import Console, ConsoleOptions, JustifyMethod, RenderResult, Segment from .jupyter import JupyterMixin from .measure import Measurement 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 delagates 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_name (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_name: str, *, theme: Union[str, SyntaxTheme] = DEFAULT_THEME, dedent: bool = False, line_numbers: bool = False, start_line: int = 1, line_range: Tuple[int, int] = None, highlight_lines: Set[int] = None, code_width: Optional[int] = None, tab_size: int = 4, word_wrap: bool = False, background_color: str = None, indent_guides: bool = False, ) -> None: self.code = code self.lexer_name = lexer_name 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: Tuple[int, int] = None, start_line: int = 1, highlight_lines: Set[int] = None, code_width: Optional[int] = None, tab_size: int = 4, word_wrap: bool = False, background_color: 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 = ( # Style(bgcolor=self.background_color) # if self.background_color is not None # else self._theme.get_background_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 def highlight(self, code: str, line_range: 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 try: lexer = get_lexer_by_name(self.lexer_name) except ClassNotFound: 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.""" 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) code = textwrap.dedent(self.code) if self.dedent else self.code code = code.expandtabs(self.tab_size) text = self.highlight(code, self.line_range) text.remove_suffix("\n") ( 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: # 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 = text.split("\n") 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") ) 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="Overide background color", ) parser.add_argument( "-x", "--lexer", default="default", dest="lexer_name", help="Lexer name", ) args = parser.parse_args() from 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_name=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)