control.py 5.2 KB

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