1from typing import Any, Callable, Dict, Iterable, List, TYPE_CHECKING, Union
2
3from .segment import ControlCode, ControlType, Segment
4
5if TYPE_CHECKING:
6    from .console import Console, ConsoleOptions, RenderResult
7
8STRIP_CONTROL_CODES = [
9    8,  # Backspace
10    11,  # Vertical tab
11    12,  # Form feed
12    13,  # Carriage return
13]
14_CONTROL_TRANSLATE = {_codepoint: None for _codepoint in STRIP_CONTROL_CODES}
15
16
17CONTROL_CODES_FORMAT: Dict[int, Callable[..., str]] = {
18    ControlType.BELL: lambda: "\x07",
19    ControlType.CARRIAGE_RETURN: lambda: "\r",
20    ControlType.HOME: lambda: "\x1b[H",
21    ControlType.CLEAR: lambda: "\x1b[2J",
22    ControlType.ENABLE_ALT_SCREEN: lambda: "\x1b[?1049h",
23    ControlType.DISABLE_ALT_SCREEN: lambda: "\x1b[?1049l",
24    ControlType.SHOW_CURSOR: lambda: "\x1b[?25h",
25    ControlType.HIDE_CURSOR: lambda: "\x1b[?25l",
26    ControlType.CURSOR_UP: lambda param: f"\x1b[{param}A",
27    ControlType.CURSOR_DOWN: lambda param: f"\x1b[{param}B",
28    ControlType.CURSOR_FORWARD: lambda param: f"\x1b[{param}C",
29    ControlType.CURSOR_BACKWARD: lambda param: f"\x1b[{param}D",
30    ControlType.CURSOR_MOVE_TO_COLUMN: lambda param: f"\x1b[{param+1}G",
31    ControlType.ERASE_IN_LINE: lambda param: f"\x1b[{param}K",
32    ControlType.CURSOR_MOVE_TO: lambda x, y: f"\x1b[{y+1};{x+1}H",
33}
34
35
36class Control:
37    """A renderable that inserts a control code (non printable but may move cursor).
38
39    Args:
40        *codes (str): Positional arguments are either a :class:`~rich.segment.ControlType` enum or a
41            tuple of ControlType and an integer parameter
42    """
43
44    __slots__ = ["segment"]
45
46    def __init__(self, *codes: Union[ControlType, ControlCode]) -> None:
47        control_codes: List[ControlCode] = [
48            (code,) if isinstance(code, ControlType) else code for code in codes
49        ]
50        _format_map = CONTROL_CODES_FORMAT
51        rendered_codes = "".join(
52            _format_map[code](*parameters) for code, *parameters in control_codes
53        )
54        self.segment = Segment(rendered_codes, None, control_codes)
55
56    @classmethod
57    def bell(cls) -> "Control":
58        """Ring the 'bell'."""
59        return cls(ControlType.BELL)
60
61    @classmethod
62    def home(cls) -> "Control":
63        """Move cursor to 'home' position."""
64        return cls(ControlType.HOME)
65
66    @classmethod
67    def move(cls, x: int = 0, y: int = 0) -> "Control":
68        """Move cursor relative to current position.
69
70        Args:
71            x (int): X offset.
72            y (int): Y offset.
73
74        Returns:
75            ~Control: Control object.
76
77        """
78
79        def get_codes() -> Iterable[ControlCode]:
80            control = ControlType
81            if x:
82                yield (
83                    control.CURSOR_FORWARD if x > 0 else control.CURSOR_BACKWARD,
84                    abs(x),
85                )
86            if y:
87                yield (
88                    control.CURSOR_DOWN if y > 0 else control.CURSOR_UP,
89                    abs(y),
90                )
91
92        control = cls(*get_codes())
93        return control
94
95    @classmethod
96    def move_to_column(cls, x: int, y: int = 0) -> "Control":
97        """Move to the given column, optionally add offset to row.
98
99        Returns:
100            x (int): absolute x (column)
101            y (int): optional y offset (row)
102
103        Returns:
104            ~Control: Control object.
105        """
106
107        return (
108            cls(
109                (ControlType.CURSOR_MOVE_TO_COLUMN, x),
110                (
111                    ControlType.CURSOR_DOWN if y > 0 else ControlType.CURSOR_UP,
112                    abs(y),
113                ),
114            )
115            if y
116            else cls((ControlType.CURSOR_MOVE_TO_COLUMN, x))
117        )
118
119    @classmethod
120    def move_to(cls, x: int, y: int) -> "Control":
121        """Move cursor to absolute position.
122
123        Args:
124            x (int): x offset (column)
125            y (int): y offset (row)
126
127        Returns:
128            ~Control: Control object.
129        """
130        return cls((ControlType.CURSOR_MOVE_TO, x, y))
131
132    @classmethod
133    def clear(cls) -> "Control":
134        """Clear the screen."""
135        return cls(ControlType.CLEAR)
136
137    @classmethod
138    def show_cursor(cls, show: bool) -> "Control":
139        """Show or hide the cursor."""
140        return cls(ControlType.SHOW_CURSOR if show else ControlType.HIDE_CURSOR)
141
142    @classmethod
143    def alt_screen(cls, enable: bool) -> "Control":
144        """Enable or disable alt screen."""
145        if enable:
146            return cls(ControlType.ENABLE_ALT_SCREEN, ControlType.HOME)
147        else:
148            return cls(ControlType.DISABLE_ALT_SCREEN)
149
150    def __str__(self) -> str:
151        return self.segment.text
152
153    def __rich_console__(
154        self, console: "Console", options: "ConsoleOptions"
155    ) -> "RenderResult":
156        if self.segment.text:
157            yield self.segment
158
159
160def strip_control_codes(
161    text: str, _translate_table: Dict[int, None] = _CONTROL_TRANSLATE
162) -> str:
163    """Remove control codes from text.
164
165    Args:
166        text (str): A string possibly contain control codes.
167
168    Returns:
169        str: String with control codes removed.
170    """
171    return text.translate(_translate_table)
172
173
174if __name__ == "__main__":  # pragma: no cover
175    print(strip_control_codes("hello\rWorld"))
176