codekingpro/portable-devtools

1import time2from typing import TYPE_CHECKING, Callable, Dict, Iterable, List, Union, Final3 4from .segment import ControlCode, ControlType, Segment5 6if TYPE_CHECKING:7    from .console import Console, ConsoleOptions, RenderResult8 9STRIP_CONTROL_CODES: Final = [10    7,  # Bell11    8,  # Backspace12    11,  # Vertical tab13    12,  # Form feed14    13,  # Carriage return15]16_CONTROL_STRIP_TRANSLATE: Final = {17    _codepoint: None for _codepoint in STRIP_CONTROL_CODES18}19 20CONTROL_ESCAPE: Final = {21    7: "\\a",22    8: "\\b",23    11: "\\v",24    12: "\\f",25    13: "\\r",26}27 28CONTROL_CODES_FORMAT: Dict[int, Callable[..., str]] = {29    ControlType.BELL: lambda: "\x07",30    ControlType.CARRIAGE_RETURN: lambda: "\r",31    ControlType.HOME: lambda: "\x1b[H",32    ControlType.CLEAR: lambda: "\x1b[2J",33    ControlType.ENABLE_ALT_SCREEN: lambda: "\x1b[?1049h",34    ControlType.DISABLE_ALT_SCREEN: lambda: "\x1b[?1049l",35    ControlType.SHOW_CURSOR: lambda: "\x1b[?25h",36    ControlType.HIDE_CURSOR: lambda: "\x1b[?25l",37    ControlType.CURSOR_UP: lambda param: f"\x1b[{param}A",38    ControlType.CURSOR_DOWN: lambda param: f"\x1b[{param}B",39    ControlType.CURSOR_FORWARD: lambda param: f"\x1b[{param}C",40    ControlType.CURSOR_BACKWARD: lambda param: f"\x1b[{param}D",41    ControlType.CURSOR_MOVE_TO_COLUMN: lambda param: f"\x1b[{param+1}G",42    ControlType.ERASE_IN_LINE: lambda param: f"\x1b[{param}K",43    ControlType.CURSOR_MOVE_TO: lambda x, y: f"\x1b[{y+1};{x+1}H",44    ControlType.SET_WINDOW_TITLE: lambda title: f"\x1b]0;{title}\x07",45}46 47 48class Control:49    """A renderable that inserts a control code (non printable but may move cursor).50 51    Args:52        *codes (str): Positional arguments are either a :class:`~rich.segment.ControlType` enum or a53            tuple of ControlType and an integer parameter54    """55 56    __slots__ = ["segment"]57 58    def __init__(self, *codes: Union[ControlType, ControlCode]) -> None:59        control_codes: List[ControlCode] = [60            (code,) if isinstance(code, ControlType) else code for code in codes61        ]62        _format_map = CONTROL_CODES_FORMAT63        rendered_codes = "".join(64            _format_map[code](*parameters) for code, *parameters in control_codes65        )66        self.segment = Segment(rendered_codes, None, control_codes)67 68    @classmethod69    def bell(cls) -> "Control":70        """Ring the 'bell'."""71        return cls(ControlType.BELL)72 73    @classmethod74    def home(cls) -> "Control":75        """Move cursor to 'home' position."""76        return cls(ControlType.HOME)77 78    @classmethod79    def move(cls, x: int = 0, y: int = 0) -> "Control":80        """Move cursor relative to current position.81 82        Args:83            x (int): X offset.84            y (int): Y offset.85 86        Returns:87            ~Control: Control object.88 89        """90 91        def get_codes() -> Iterable[ControlCode]:92            control = ControlType93            if x:94                yield (95                    control.CURSOR_FORWARD if x > 0 else control.CURSOR_BACKWARD,96                    abs(x),97                )98            if y:99                yield (100                    control.CURSOR_DOWN if y > 0 else control.CURSOR_UP,101                    abs(y),102                )103 104        control = cls(*get_codes())105        return control106 107    @classmethod108    def move_to_column(cls, x: int, y: int = 0) -> "Control":109        """Move to the given column, optionally add offset to row.110 111        Returns:112            x (int): absolute x (column)113            y (int): optional y offset (row)114 115        Returns:116            ~Control: Control object.117        """118 119        return (120            cls(121                (ControlType.CURSOR_MOVE_TO_COLUMN, x),122                (123                    ControlType.CURSOR_DOWN if y > 0 else ControlType.CURSOR_UP,124                    abs(y),125                ),126            )127            if y128            else cls((ControlType.CURSOR_MOVE_TO_COLUMN, x))129        )130 131    @classmethod132    def move_to(cls, x: int, y: int) -> "Control":133        """Move cursor to absolute position.134 135        Args:136            x (int): x offset (column)137            y (int): y offset (row)138 139        Returns:140            ~Control: Control object.141        """142        return cls((ControlType.CURSOR_MOVE_TO, x, y))143 144    @classmethod145    def clear(cls) -> "Control":146        """Clear the screen."""147        return cls(ControlType.CLEAR)148 149    @classmethod150    def show_cursor(cls, show: bool) -> "Control":151        """Show or hide the cursor."""152        return cls(ControlType.SHOW_CURSOR if show else ControlType.HIDE_CURSOR)153 154    @classmethod155    def alt_screen(cls, enable: bool) -> "Control":156        """Enable or disable alt screen."""157        if enable:158            return cls(ControlType.ENABLE_ALT_SCREEN, ControlType.HOME)159        else:160            return cls(ControlType.DISABLE_ALT_SCREEN)161 162    @classmethod163    def title(cls, title: str) -> "Control":164        """Set the terminal window title165 166        Args:167            title (str): The new terminal window title168        """169        return cls((ControlType.SET_WINDOW_TITLE, title))170 171    def __str__(self) -> str:172        return self.segment.text173 174    def __rich_console__(175        self, console: "Console", options: "ConsoleOptions"176    ) -> "RenderResult":177        if self.segment.text:178            yield self.segment179 180 181def strip_control_codes(182    text: str, _translate_table: Dict[int, None] = _CONTROL_STRIP_TRANSLATE183) -> str:184    """Remove control codes from text.185 186    Args:187        text (str): A string possibly contain control codes.188 189    Returns:190        str: String with control codes removed.191    """192    return text.translate(_translate_table)193 194 195def escape_control_codes(196    text: str,197    _translate_table: Dict[int, str] = CONTROL_ESCAPE,198) -> str:199    """Replace control codes with their "escaped" equivalent in the given text.200    (e.g. "\b" becomes "\\b")201 202    Args:203        text (str): A string possibly containing control codes.204 205    Returns:206        str: String with control codes replaced with their escaped version.207    """208    return text.translate(_translate_table)209 210 211if __name__ == "__main__":  # pragma: no cover212    from rich.console import Console213 214    console = Console()215    console.print("Look at the title of your terminal window ^")216    # console.print(Control((ControlType.SET_WINDOW_TITLE, "Hello, world!")))217    for i in range(10):218        console.set_window_title("🚀 Loading" + "." * i)219        time.sleep(0.5)220