control.py 6.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225
  1. import sys
  2. import time
  3. from typing import TYPE_CHECKING, Callable, Dict, Iterable, List, Union
  4. if sys.version_info >= (3, 8):
  5. from typing import Final
  6. else:
  7. from pip._vendor.typing_extensions import Final # pragma: no cover
  8. from .segment import ControlCode, ControlType, Segment
  9. if TYPE_CHECKING:
  10. from .console import Console, ConsoleOptions, RenderResult
  11. STRIP_CONTROL_CODES: Final = [
  12. 7, # Bell
  13. 8, # Backspace
  14. 11, # Vertical tab
  15. 12, # Form feed
  16. 13, # Carriage return
  17. ]
  18. _CONTROL_STRIP_TRANSLATE: Final = {
  19. _codepoint: None for _codepoint in STRIP_CONTROL_CODES
  20. }
  21. CONTROL_ESCAPE: Final = {
  22. 7: "\\a",
  23. 8: "\\b",
  24. 11: "\\v",
  25. 12: "\\f",
  26. 13: "\\r",
  27. }
  28. CONTROL_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. class Control:
  47. """A renderable that inserts a control code (non printable but may move cursor).
  48. Args:
  49. *codes (str): Positional arguments are either a :class:`~rich.segment.ControlType` enum or a
  50. tuple of ControlType and an integer parameter
  51. """
  52. __slots__ = ["segment"]
  53. def __init__(self, *codes: Union[ControlType, ControlCode]) -> None:
  54. control_codes: List[ControlCode] = [
  55. (code,) if isinstance(code, ControlType) else code for code in codes
  56. ]
  57. _format_map = CONTROL_CODES_FORMAT
  58. rendered_codes = "".join(
  59. _format_map[code](*parameters) for code, *parameters in control_codes
  60. )
  61. self.segment = Segment(rendered_codes, None, control_codes)
  62. @classmethod
  63. def bell(cls) -> "Control":
  64. """Ring the 'bell'."""
  65. return cls(ControlType.BELL)
  66. @classmethod
  67. def home(cls) -> "Control":
  68. """Move cursor to 'home' position."""
  69. return cls(ControlType.HOME)
  70. @classmethod
  71. def move(cls, x: int = 0, y: int = 0) -> "Control":
  72. """Move cursor relative to current position.
  73. Args:
  74. x (int): X offset.
  75. y (int): Y offset.
  76. Returns:
  77. ~Control: Control object.
  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. control = cls(*get_codes())
  92. return control
  93. @classmethod
  94. def move_to_column(cls, x: int, y: int = 0) -> "Control":
  95. """Move to the given column, optionally add offset to row.
  96. Returns:
  97. x (int): absolute x (column)
  98. y (int): optional y offset (row)
  99. Returns:
  100. ~Control: Control object.
  101. """
  102. return (
  103. cls(
  104. (ControlType.CURSOR_MOVE_TO_COLUMN, x),
  105. (
  106. ControlType.CURSOR_DOWN if y > 0 else ControlType.CURSOR_UP,
  107. abs(y),
  108. ),
  109. )
  110. if y
  111. else cls((ControlType.CURSOR_MOVE_TO_COLUMN, x))
  112. )
  113. @classmethod
  114. def move_to(cls, x: int, y: int) -> "Control":
  115. """Move cursor to absolute position.
  116. Args:
  117. x (int): x offset (column)
  118. y (int): y offset (row)
  119. Returns:
  120. ~Control: Control object.
  121. """
  122. return cls((ControlType.CURSOR_MOVE_TO, x, y))
  123. @classmethod
  124. def clear(cls) -> "Control":
  125. """Clear the screen."""
  126. return cls(ControlType.CLEAR)
  127. @classmethod
  128. def show_cursor(cls, show: bool) -> "Control":
  129. """Show or hide the cursor."""
  130. return cls(ControlType.SHOW_CURSOR if show else ControlType.HIDE_CURSOR)
  131. @classmethod
  132. def alt_screen(cls, enable: bool) -> "Control":
  133. """Enable or disable alt screen."""
  134. if enable:
  135. return cls(ControlType.ENABLE_ALT_SCREEN, ControlType.HOME)
  136. else:
  137. return cls(ControlType.DISABLE_ALT_SCREEN)
  138. @classmethod
  139. def title(cls, title: str) -> "Control":
  140. """Set the terminal window title
  141. Args:
  142. title (str): The new terminal window title
  143. """
  144. return cls((ControlType.SET_WINDOW_TITLE, title))
  145. def __str__(self) -> str:
  146. return self.segment.text
  147. def __rich_console__(
  148. self, console: "Console", options: "ConsoleOptions"
  149. ) -> "RenderResult":
  150. if self.segment.text:
  151. yield self.segment
  152. def strip_control_codes(
  153. text: str, _translate_table: Dict[int, None] = _CONTROL_STRIP_TRANSLATE
  154. ) -> str:
  155. """Remove control codes from text.
  156. Args:
  157. text (str): A string possibly contain control codes.
  158. Returns:
  159. str: String with control codes removed.
  160. """
  161. return text.translate(_translate_table)
  162. def escape_control_codes(
  163. text: str,
  164. _translate_table: Dict[int, str] = CONTROL_ESCAPE,
  165. ) -> str:
  166. """Replace control codes with their "escaped" equivalent in the given text.
  167. (e.g. "\b" becomes "\\b")
  168. Args:
  169. text (str): A string possibly containing control codes.
  170. Returns:
  171. str: String with control codes replaced with their escaped version.
  172. """
  173. return text.translate(_translate_table)
  174. if __name__ == "__main__": # pragma: no cover
  175. from pip._vendor.rich.console import Console
  176. console = Console()
  177. console.print("Look at the title of your terminal window ^")
  178. # console.print(Control((ControlType.SET_WINDOW_TITLE, "Hello, world!")))
  179. for i in range(10):
  180. console.set_window_title("🚀 Loading" + "." * i)
  181. time.sleep(0.5)