1#!/usr/local/bin/python3.8
2# vim:fileencoding=utf-8
3# License: GPL v3 Copyright: 2018, Kovid Goyal <kovid at kovidgoyal.net>
4
5import sys
6from contextlib import contextmanager
7from functools import wraps
8from enum import Enum
9from typing import (
10    IO, Any, Callable, Dict, Generator, Optional, Tuple, TypeVar, Union
11)
12
13from kitty.rgb import Color, color_as_sharp, to_color
14from kitty.typing import GraphicsCommandType, HandlerType, ScreenSize
15
16from .operations_stub import CMD
17
18GraphicsCommandType, ScreenSize  # needed for stub generation
19S7C1T = '\033 F'
20SAVE_CURSOR = '\0337'
21RESTORE_CURSOR = '\0338'
22SAVE_PRIVATE_MODE_VALUES = '\033[?s'
23RESTORE_PRIVATE_MODE_VALUES = '\033[?r'
24SAVE_COLORS = '\033[#P'
25RESTORE_COLORS = '\033[#Q'
26F = TypeVar('F')
27all_cmds: Dict[str, Callable] = {}
28
29
30class Mode(Enum):
31    LNM = 20, ''
32    IRM = 4, ''
33    DECKM = 1, '?'
34    DECSCNM = 5, '?'
35    DECOM = 6, '?'
36    DECAWM = 7, '?'
37    DECARM = 8, '?'
38    DECTCEM = 25, '?'
39    MOUSE_BUTTON_TRACKING = 1000, '?'
40    MOUSE_MOTION_TRACKING = 1002, '?'
41    MOUSE_MOVE_TRACKING = 1003, '?'
42    FOCUS_TRACKING = 1004, '?'
43    MOUSE_UTF8_MODE = 1005, '?'
44    MOUSE_SGR_MODE = 1006, '?'
45    MOUSE_URXVT_MODE = 1015, '?'
46    ALTERNATE_SCREEN = 1049, '?'
47    BRACKETED_PASTE = 2004, '?'
48    PENDING_UPDATE = 2026, '?'
49
50
51def cmd(f: F) -> F:
52    all_cmds[f.__name__] = f  # type: ignore
53    return f
54
55
56@cmd
57def set_mode(which: Mode) -> str:
58    num, private = which.value
59    return '\033[{}{}h'.format(private, num)
60
61
62@cmd
63def reset_mode(which: Mode) -> str:
64    num, private = which.value
65    return '\033[{}{}l'.format(private, num)
66
67
68@cmd
69def clear_screen() -> str:
70    return '\033[H\033[2J'
71
72
73@cmd
74def clear_to_end_of_screen() -> str:
75    return '\033[J'
76
77
78@cmd
79def clear_to_eol() -> str:
80    return '\033[K'
81
82
83@cmd
84def reset_terminal() -> str:
85    return '\033]\033\\\033c'
86
87
88@cmd
89def bell() -> str:
90    return '\a'
91
92
93@cmd
94def beep() -> str:
95    return '\a'
96
97
98@cmd
99def set_window_title(value: str) -> str:
100    return '\033]2;' + value.replace('\033', '').replace('\x9c', '') + '\033\\'
101
102
103@cmd
104def set_line_wrapping(yes_or_no: bool) -> str:
105    return set_mode(Mode.DECAWM) if yes_or_no else reset_mode(Mode.DECAWM)
106
107
108@cmd
109def set_cursor_visible(yes_or_no: bool) -> str:
110    return set_mode(Mode.DECTCEM) if yes_or_no else reset_mode(Mode.DECTCEM)
111
112
113@cmd
114def set_cursor_position(x: int = 0, y: int = 0) -> str:  # (0, 0) is top left
115    return '\033[{};{}H'.format(y + 1, x + 1)
116
117
118@cmd
119def move_cursor_by(amt: int, direction: str) -> str:
120    suffix = {'up': 'A', 'down': 'B', 'right': 'C', 'left': 'D'}[direction]
121    return f'\033[{amt}{suffix}'
122
123
124@cmd
125def set_cursor_shape(shape: str = 'block', blink: bool = True) -> str:
126    val = {'block': 1, 'underline': 3, 'bar': 5}.get(shape, 1)
127    if not blink:
128        val += 1
129    return '\033[{} q'.format(val)
130
131
132@cmd
133def set_scrolling_region(screen_size: Optional['ScreenSize'] = None, top: Optional[int] = None, bottom: Optional[int] = None) -> str:
134    if screen_size is None:
135        return '\033[r'
136    if top is None:
137        top = 0
138    if bottom is None:
139        bottom = screen_size.rows - 1
140    if bottom < 0:
141        bottom = screen_size.rows - 1 + bottom
142    else:
143        bottom += 1
144    return '\033[{};{}r'.format(top + 1, bottom + 1)
145
146
147@cmd
148def scroll_screen(amt: int = 1) -> str:
149    return '\033[' + str(abs(amt)) + ('T' if amt < 0 else 'S')
150
151
152STANDARD_COLORS = {name: i for i, name in enumerate(
153    'black red green yellow blue magenta cyan gray'.split())}
154STANDARD_COLORS['white'] = STANDARD_COLORS['gray']
155UNDERLINE_STYLES = {name: i + 1 for i, name in enumerate(
156    'straight double curly'.split())}
157
158
159ColorSpec = Union[int, str, Tuple[int, int, int]]
160
161
162def color_code(color: ColorSpec, intense: bool = False, base: int = 30) -> str:
163    if isinstance(color, str):
164        e = str((base + 60 if intense else base) + STANDARD_COLORS[color])
165    elif isinstance(color, int):
166        e = '{}:5:{}'.format(base + 8, max(0, min(color, 255)))
167    else:
168        e = '{}:2:{}:{}:{}'.format(base + 8, *color)
169    return e
170
171
172@cmd
173def sgr(*parts: str) -> str:
174    return '\033[{}m'.format(';'.join(parts))
175
176
177@cmd
178def colored(
179    text: str,
180    color: ColorSpec,
181    intense: bool = False,
182    reset_to: Optional[ColorSpec] = None,
183    reset_to_intense: bool = False
184) -> str:
185    e = color_code(color, intense)
186    return '\033[{}m{}\033[{}m'.format(e, text, 39 if reset_to is None else color_code(reset_to, reset_to_intense))
187
188
189@cmd
190def faint(text: str) -> str:
191    return colored(text, 'black', True)
192
193
194@cmd
195def styled(
196    text: str,
197    fg: Optional[ColorSpec] = None,
198    bg: Optional[ColorSpec] = None,
199    fg_intense: bool = False,
200    bg_intense: bool = False,
201    italic: Optional[bool] = None,
202    bold: Optional[bool] = None,
203    underline: Optional[str] = None,
204    underline_color: Optional[ColorSpec] = None,
205    reverse: Optional[bool] = None
206) -> str:
207    start, end = [], []
208    if fg is not None:
209        start.append(color_code(fg, fg_intense))
210        end.append('39')
211    if bg is not None:
212        start.append(color_code(bg, bg_intense, 40))
213        end.append('49')
214    if underline_color is not None:
215        if isinstance(underline_color, str):
216            underline_color = STANDARD_COLORS[underline_color]
217        start.append(color_code(underline_color, base=50))
218        end.append('59')
219    if underline is not None:
220        start.append('4:{}'.format(UNDERLINE_STYLES[underline]))
221        end.append('4:0')
222    if italic is not None:
223        s, e = (start, end) if italic else (end, start)
224        s.append('3')
225        e.append('23')
226    if bold is not None:
227        s, e = (start, end) if bold else (end, start)
228        s.append('1')
229        e.append('22')
230    if reverse is not None:
231        s, e = (start, end) if reverse else (end, start)
232        s.append('7')
233        e.append('27')
234    if not start:
235        return text
236    return '\033[{}m{}\033[{}m'.format(';'.join(start), text, ';'.join(end))
237
238
239def serialize_gr_command(cmd: Dict[str, Union[int, str]], payload: Optional[bytes] = None) -> bytes:
240    from .images import GraphicsCommand
241    gc = GraphicsCommand()
242    for k, v in cmd.items():
243        setattr(gc, k, v)
244    return gc.serialize(payload or b'')
245
246
247@cmd
248def gr_command(cmd: Union[Dict, 'GraphicsCommandType'], payload: Optional[bytes] = None) -> str:
249    if isinstance(cmd, dict):
250        raw = serialize_gr_command(cmd, payload)
251    else:
252        raw = cmd.serialize(payload or b'')
253    return raw.decode('ascii')
254
255
256@cmd
257def clear_images_on_screen(delete_data: bool = False) -> str:
258    from .images import GraphicsCommand
259    gc = GraphicsCommand()
260    gc.a = 'd'
261    gc.d = 'A' if delete_data else 'a'
262    return gc.serialize().decode('ascii')
263
264
265def init_state(alternate_screen: bool = True) -> str:
266    ans = (
267        S7C1T + SAVE_CURSOR + SAVE_PRIVATE_MODE_VALUES + reset_mode(Mode.LNM) +
268        reset_mode(Mode.IRM) + reset_mode(Mode.DECKM) + reset_mode(Mode.DECSCNM) +
269        set_mode(Mode.DECARM) + set_mode(Mode.DECAWM) +
270        set_mode(Mode.DECTCEM) + reset_mode(Mode.MOUSE_BUTTON_TRACKING) +
271        reset_mode(Mode.MOUSE_MOTION_TRACKING) + reset_mode(Mode.MOUSE_MOVE_TRACKING) +
272        reset_mode(Mode.FOCUS_TRACKING) + reset_mode(Mode.MOUSE_UTF8_MODE) +
273        reset_mode(Mode.MOUSE_SGR_MODE) + reset_mode(Mode.MOUSE_UTF8_MODE) +
274        set_mode(Mode.BRACKETED_PASTE) + SAVE_COLORS +
275        '\033[*x'  # reset DECSACE to default region select
276    )
277    if alternate_screen:
278        ans += set_mode(Mode.ALTERNATE_SCREEN) + reset_mode(Mode.DECOM)
279        ans += clear_screen()
280    ans += '\033[>31u'  # extended keyboard mode
281    return ans
282
283
284def reset_state(normal_screen: bool = True) -> str:
285    ans = ''
286    ans += '\033[<u'  # restore keyboard mode
287    if normal_screen:
288        ans += reset_mode(Mode.ALTERNATE_SCREEN)
289    ans += RESTORE_PRIVATE_MODE_VALUES
290    ans += RESTORE_CURSOR
291    ans += RESTORE_COLORS
292    return ans
293
294
295@contextmanager
296def pending_update(write: Callable[[str], None]) -> Generator[None, None, None]:
297    write(set_mode(Mode.PENDING_UPDATE))
298    yield
299    write(reset_mode(Mode.PENDING_UPDATE))
300
301
302@contextmanager
303def cursor(write: Callable[[str], None]) -> Generator[None, None, None]:
304    write(SAVE_CURSOR)
305    yield
306    write(RESTORE_CURSOR)
307
308
309@contextmanager
310def alternate_screen(f: Optional[IO[str]] = None) -> Generator[None, None, None]:
311    f = f or sys.stdout
312    print(set_mode(Mode.ALTERNATE_SCREEN), end='', file=f)
313    yield
314    print(reset_mode(Mode.ALTERNATE_SCREEN), end='', file=f)
315
316
317@contextmanager
318def raw_mode(fd: Optional[int] = None) -> Generator[None, None, None]:
319    import tty
320    import termios
321    if fd is None:
322        fd = sys.stdin.fileno()
323    old = termios.tcgetattr(fd)
324    try:
325        tty.setraw(fd)
326        yield
327    finally:
328        termios.tcsetattr(fd, termios.TCSADRAIN, old)
329
330
331@cmd
332def set_default_colors(
333    fg: Optional[Union[Color, str]] = None,
334    bg: Optional[Union[Color, str]] = None,
335    cursor: Optional[Union[Color, str]] = None,
336    select_bg: Optional[Union[Color, str]] = None,
337    select_fg: Optional[Union[Color, str]] = None
338) -> str:
339    ans = ''
340
341    def item(which: Optional[Union[Color, str]], num: int) -> None:
342        nonlocal ans
343        if which is None:
344            ans += '\x1b]1{}\x1b\\'.format(num)
345        else:
346            if isinstance(which, Color):
347                q = color_as_sharp(which)
348            else:
349                x = to_color(which)
350                assert x is not None
351                q = color_as_sharp(x)
352            ans += '\x1b]{};{}\x1b\\'.format(num, q)
353
354    item(fg, 10)
355    item(bg, 11)
356    item(cursor, 12)
357    item(select_bg, 17)
358    item(select_fg, 19)
359    return ans
360
361
362@cmd
363def save_colors() -> str:
364    return '\x1b[#P'
365
366
367@cmd
368def restore_colors() -> str:
369    return '\x1b[#Q'
370
371
372@cmd
373def write_to_clipboard(data: Union[str, bytes], use_primary: bool = False) -> str:
374    from base64 import standard_b64encode
375    fmt = 'p' if use_primary else 'c'
376    if isinstance(data, str):
377        data = data.encode('utf-8')
378    payload = standard_b64encode(data).decode('ascii')
379    return f'\x1b]52;{fmt};{payload}\a'
380
381
382@cmd
383def request_from_clipboard(use_primary: bool = False) -> str:
384    return '\x1b]52;{};?\a'.format('p' if use_primary else 'c')
385
386
387# Boilerplate to make operations available via Handler.cmd  {{{
388
389
390def writer(handler: HandlerType, func: Callable) -> Callable:
391    @wraps(func)
392    def f(*a: Any, **kw: Any) -> None:
393        handler.write(func(*a, **kw))
394    return f
395
396
397def commander(handler: HandlerType) -> CMD:
398    ans = CMD()
399    for name, func in all_cmds.items():
400        setattr(ans, name, writer(handler, func))
401    return ans
402
403
404def func_sig(func: Callable) -> Generator[str, None, None]:
405    import inspect
406    import re
407    s = inspect.signature(func)
408    for val in s.parameters.values():
409        yield re.sub(r'ForwardRef\([\'"](\w+?)[\'"]\)', r'\1', str(val).replace('NoneType', 'None'))
410
411
412def as_type_stub() -> str:
413    ans = [
414        'from typing import *  # noqa',
415        'from kitty.typing import GraphicsCommandType, ScreenSize',
416        'from kitty.rgb import Color',
417        'import kitty.rgb',
418        'import kittens.tui.operations',
419    ]
420    methods = []
421    for name, func in all_cmds.items():
422        args = ', '.join(func_sig(func))
423        if args:
424            args = ', ' + args
425        methods.append('    def {}(self{}) -> str: pass'.format(name, args))
426    ans += ['', '', 'class CMD:'] + methods
427
428    return '\n'.join(ans) + '\n\n\n'
429# }}}
430