summaryrefslogtreecommitdiffstats
path: root/mycli/clistyle.py
blob: b0ac99221857316b4d39429b6742eccc10dce527 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
import logging

import pygments.styles
from pygments.token import string_to_tokentype, Token
from pygments.style import Style as PygmentsStyle
from pygments.util import ClassNotFound
from prompt_toolkit.styles.pygments import style_from_pygments_cls
from prompt_toolkit.styles import merge_styles, Style

logger = logging.getLogger(__name__)

# map Pygments tokens (ptk 1.0) to class names (ptk 2.0).
TOKEN_TO_PROMPT_STYLE = {
    Token.Menu.Completions.Completion.Current: 'completion-menu.completion.current',
    Token.Menu.Completions.Completion: 'completion-menu.completion',
    Token.Menu.Completions.Meta.Current: 'completion-menu.meta.completion.current',
    Token.Menu.Completions.Meta: 'completion-menu.meta.completion',
    Token.Menu.Completions.MultiColumnMeta: 'completion-menu.multi-column-meta',
    Token.Menu.Completions.ProgressButton: 'scrollbar.arrow',  # best guess
    Token.Menu.Completions.ProgressBar: 'scrollbar',  # best guess
    Token.SelectedText: 'selected',
    Token.SearchMatch: 'search',
    Token.SearchMatch.Current: 'search.current',
    Token.Toolbar: 'bottom-toolbar',
    Token.Toolbar.Off: 'bottom-toolbar.off',
    Token.Toolbar.On: 'bottom-toolbar.on',
    Token.Toolbar.Search: 'search-toolbar',
    Token.Toolbar.Search.Text: 'search-toolbar.text',
    Token.Toolbar.System: 'system-toolbar',
    Token.Toolbar.Arg: 'arg-toolbar',
    Token.Toolbar.Arg.Text: 'arg-toolbar.text',
    Token.Toolbar.Transaction.Valid: 'bottom-toolbar.transaction.valid',
    Token.Toolbar.Transaction.Failed: 'bottom-toolbar.transaction.failed',
    Token.Output.Header: 'output.header',
    Token.Output.OddRow: 'output.odd-row',
    Token.Output.EvenRow: 'output.even-row',
    Token.Output.Null: 'output.null',
    Token.Prompt: 'prompt',
    Token.Continuation: 'continuation',
}

# reverse dict for cli_helpers, because they still expect Pygments tokens.
PROMPT_STYLE_TO_TOKEN = {
    v: k for k, v in TOKEN_TO_PROMPT_STYLE.items()
}

# all tokens that the Pygments MySQL lexer can produce
OVERRIDE_STYLE_TO_TOKEN = {
    'sql.comment': Token.Comment,
    'sql.comment.multi-line': Token.Comment.Multiline,
    'sql.comment.single-line': Token.Comment.Single,
    'sql.comment.optimizer-hint': Token.Comment.Special,
    'sql.escape': Token.Error,
    'sql.keyword': Token.Keyword,
    'sql.datatype': Token.Keyword.Type,
    'sql.literal': Token.Literal,
    'sql.literal.date': Token.Literal.Date,
    'sql.symbol': Token.Name,
    'sql.quoted-schema-object': Token.Name.Quoted,
    'sql.quoted-schema-object.escape': Token.Name.Quoted.Escape,
    'sql.constant': Token.Name.Constant,
    'sql.function': Token.Name.Function,
    'sql.variable': Token.Name.Variable,
    'sql.number': Token.Number,
    'sql.number.binary': Token.Number.Bin,
    'sql.number.float': Token.Number.Float,
    'sql.number.hex': Token.Number.Hex,
    'sql.number.integer': Token.Number.Integer,
    'sql.operator': Token.Operator,
    'sql.punctuation': Token.Punctuation,
    'sql.string': Token.String,
    'sql.string.double-quouted': Token.String.Double,
    'sql.string.escape': Token.String.Escape,
    'sql.string.single-quoted': Token.String.Single,
    'sql.whitespace': Token.Text,
}

def parse_pygments_style(token_name, style_object, style_dict):
    """Parse token type and style string.

    :param token_name: str name of Pygments token. Example: "Token.String"
    :param style_object: pygments.style.Style instance to use as base
    :param style_dict: dict of token names and their styles, customized to this cli

    """
    token_type = string_to_tokentype(token_name)
    try:
        other_token_type = string_to_tokentype(style_dict[token_name])
        return token_type, style_object.styles[other_token_type]
    except AttributeError as err:
        return token_type, style_dict[token_name]


def style_factory(name, cli_style):
    try:
        style = pygments.styles.get_style_by_name(name)
    except ClassNotFound:
        style = pygments.styles.get_style_by_name('native')

    prompt_styles = []
    # prompt-toolkit used pygments tokens for styling before, switched to style
    # names in 2.0. Convert old token types to new style names, for backwards compatibility.
    for token in cli_style:
        if token.startswith('Token.'):
            # treat as pygments token (1.0)
            token_type, style_value = parse_pygments_style(
                token, style, cli_style)
            if token_type in TOKEN_TO_PROMPT_STYLE:
                prompt_style = TOKEN_TO_PROMPT_STYLE[token_type]
                prompt_styles.append((prompt_style, style_value))
            else:
                # we don't want to support tokens anymore
                logger.error('Unhandled style / class name: %s', token)
        else:
            # treat as prompt style name (2.0). See default style names here:
            # https://github.com/jonathanslenders/python-prompt-toolkit/blob/master/prompt_toolkit/styles/defaults.py
            prompt_styles.append((token, cli_style[token]))

    override_style = Style([('bottom-toolbar', 'noreverse')])
    return merge_styles([
        style_from_pygments_cls(style),
        override_style,
        Style(prompt_styles)
    ])


def style_factory_output(name, cli_style):
    try:
        style = pygments.styles.get_style_by_name(name).styles
    except ClassNotFound:
        style = pygments.styles.get_style_by_name('native').styles

    for token in cli_style:
        if token.startswith('Token.'):
            token_type, style_value = parse_pygments_style(
                token, style, cli_style)
            style.update({token_type: style_value})
        elif token in PROMPT_STYLE_TO_TOKEN:
            token_type = PROMPT_STYLE_TO_TOKEN[token]
            style.update({token_type: cli_style[token]})
        elif token in OVERRIDE_STYLE_TO_TOKEN:
            token_type = OVERRIDE_STYLE_TO_TOKEN[token]
            style.update({token_type: cli_style[token]})
        else:
            # TODO: cli helpers will have to switch to ptk.Style
            logger.error('Unhandled style / class name: %s', token)

    class OutputStyle(PygmentsStyle):
        default_style = ""
        styles = style

    return OutputStyle