123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175 |
- from typing import Any, Callable, Dict, Iterable, List, TYPE_CHECKING, Union
- from .segment import ControlCode, ControlType, Segment
- if TYPE_CHECKING:
- from .console import Console, ConsoleOptions, RenderResult
- STRIP_CONTROL_CODES = [
- 8, # Backspace
- 11, # Vertical tab
- 12, # Form feed
- 13, # Carriage return
- ]
- _CONTROL_TRANSLATE = {_codepoint: None for _codepoint in STRIP_CONTROL_CODES}
- CONTROL_CODES_FORMAT: Dict[int, Callable[..., str]] = {
- ControlType.BELL: lambda: "\x07",
- ControlType.CARRIAGE_RETURN: lambda: "\r",
- ControlType.HOME: lambda: "\x1b[H",
- ControlType.CLEAR: lambda: "\x1b[2J",
- ControlType.ENABLE_ALT_SCREEN: lambda: "\x1b[?1049h",
- ControlType.DISABLE_ALT_SCREEN: lambda: "\x1b[?1049l",
- ControlType.SHOW_CURSOR: lambda: "\x1b[?25h",
- ControlType.HIDE_CURSOR: lambda: "\x1b[?25l",
- ControlType.CURSOR_UP: lambda param: f"\x1b[{param}A",
- ControlType.CURSOR_DOWN: lambda param: f"\x1b[{param}B",
- ControlType.CURSOR_FORWARD: lambda param: f"\x1b[{param}C",
- ControlType.CURSOR_BACKWARD: lambda param: f"\x1b[{param}D",
- ControlType.CURSOR_MOVE_TO_COLUMN: lambda param: f"\x1b[{param+1}G",
- ControlType.ERASE_IN_LINE: lambda param: f"\x1b[{param}K",
- ControlType.CURSOR_MOVE_TO: lambda x, y: f"\x1b[{y+1};{x+1}H",
- }
- class Control:
- """A renderable that inserts a control code (non printable but may move cursor).
- Args:
- *codes (str): Positional arguments are either a :class:`~rich.segment.ControlType` enum or a
- tuple of ControlType and an integer parameter
- """
- __slots__ = ["segment"]
- def __init__(self, *codes: Union[ControlType, ControlCode]) -> None:
- control_codes: List[ControlCode] = [
- (code,) if isinstance(code, ControlType) else code for code in codes
- ]
- _format_map = CONTROL_CODES_FORMAT
- rendered_codes = "".join(
- _format_map[code](*parameters) for code, *parameters in control_codes
- )
- self.segment = Segment(rendered_codes, None, control_codes)
- @classmethod
- def bell(cls) -> "Control":
- """Ring the 'bell'."""
- return cls(ControlType.BELL)
- @classmethod
- def home(cls) -> "Control":
- """Move cursor to 'home' position."""
- return cls(ControlType.HOME)
- @classmethod
- def move(cls, x: int = 0, y: int = 0) -> "Control":
- """Move cursor relative to current position.
- Args:
- x (int): X offset.
- y (int): Y offset.
- Returns:
- ~Control: Control object.
- """
- def get_codes() -> Iterable[ControlCode]:
- control = ControlType
- if x:
- yield (
- control.CURSOR_FORWARD if x > 0 else control.CURSOR_BACKWARD,
- abs(x),
- )
- if y:
- yield (
- control.CURSOR_DOWN if y > 0 else control.CURSOR_UP,
- abs(y),
- )
- control = cls(*get_codes())
- return control
- @classmethod
- def move_to_column(cls, x: int, y: int = 0) -> "Control":
- """Move to the given column, optionally add offset to row.
- Returns:
- x (int): absolute x (column)
- y (int): optional y offset (row)
- Returns:
- ~Control: Control object.
- """
- return (
- cls(
- (ControlType.CURSOR_MOVE_TO_COLUMN, x),
- (
- ControlType.CURSOR_DOWN if y > 0 else ControlType.CURSOR_UP,
- abs(y),
- ),
- )
- if y
- else cls((ControlType.CURSOR_MOVE_TO_COLUMN, x))
- )
- @classmethod
- def move_to(cls, x: int, y: int) -> "Control":
- """Move cursor to absolute position.
- Args:
- x (int): x offset (column)
- y (int): y offset (row)
- Returns:
- ~Control: Control object.
- """
- return cls((ControlType.CURSOR_MOVE_TO, x, y))
- @classmethod
- def clear(cls) -> "Control":
- """Clear the screen."""
- return cls(ControlType.CLEAR)
- @classmethod
- def show_cursor(cls, show: bool) -> "Control":
- """Show or hide the cursor."""
- return cls(ControlType.SHOW_CURSOR if show else ControlType.HIDE_CURSOR)
- @classmethod
- def alt_screen(cls, enable: bool) -> "Control":
- """Enable or disable alt screen."""
- if enable:
- return cls(ControlType.ENABLE_ALT_SCREEN, ControlType.HOME)
- else:
- return cls(ControlType.DISABLE_ALT_SCREEN)
- def __str__(self) -> str:
- return self.segment.text
- def __rich_console__(
- self, console: "Console", options: "ConsoleOptions"
- ) -> "RenderResult":
- if self.segment.text:
- yield self.segment
- def strip_control_codes(
- text: str, _translate_table: Dict[int, None] = _CONTROL_TRANSLATE
- ) -> str:
- """Remove control codes from text.
- Args:
- text (str): A string possibly contain control codes.
- Returns:
- str: String with control codes removed.
- """
- return text.translate(_translate_table)
- if __name__ == "__main__": # pragma: no cover
- print(strip_control_codes("hello\rWorld"))
|