capture

CHANGELOG.md

@@ -5,7 +5,11 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
-## [7.0.1] - unreleased
+## [7.1.0] - unreleased
+
+### Added
+
+- Added Console.capture method
 
 ### Changed
 

README.md

@@ -54,6 +54,15 @@ Rich can be installed in the Python REPL, so that any data structures will be pr
 
 ![REPL](https://github.com/willmcgugan/rich/raw/master/imgs/repl.png)
 
+## Rich Inspect
+
+Rich has an [inspect](https://rich.readthedocs.io/en/latest/reference/init.html?highlight=inspect#rich.inspect) function which can produce a report on any Python object, such as class, instance, or builtin.
+
+```python
+>>> from rich import inspect
+>>> inspect(str, methods=True)
+```
+
 ## Using the Console
 
 For more control over rich terminal content, import and construct a [Console](https://rich.readthedocs.io/en/latest/reference/console.html#rich.console.Console) object.

docs/source/console.rst

@@ -42,48 +42,6 @@ You can set ``color_system`` to one of the following values:
     Be careful when setting a color system, if you set a higher color system than your terminal supports, your text may be unreadable.
 
 
-File output
------------
-
-The Console object will write to standard output (i.e. the terminal). You can also tell the Console object to write to another file by setting the ``file`` argument on the constructor -- which should be a file-like object opened for writing text. One use of this capability is to create a Console for writing to standard error by setting file to ``sys.stderr``. Here's an example::
-
-    import sys
-    from rich.console import Console
-    error_console = Console(file=sys.stderr)
-    error_console.print("[bold red]This is an error!")
-
-
-Capturing output
-----------------
-
-There may be situations where you want to capture the output from a Console rather than writing it directly to the terminal. You can do this by setting the ``file`` argument to a :py:class:`io.StringIO` instance. Here's an example::
-
-    from io import StringIO
-    from rich.console import Console
-    console = Console(file=StringIO())
-    console.print("[bold red]Hello[/] World")
-    str_output = console.file.getvalue()
-
-You may also want to set ``force_terminal=True`` on the Console constructor if you want control codes for colour and style in the resulting string.
-
-Terminal detection
-------------------
-
-If Rich detects that it is not writing to a terminal it will strip control codes from the output. If you want to write control codes to a regular file then set ``force_terminal=True`` on the constructor.
-
-Letting Rich auto-detect terminals is useful as it will write plain text when you pipe output to a file or other application.
-
-
-Environment variables
----------------------
-
-Rich respects some standard environment variables.
-
-Setting the environment variable ``TERM`` to ``"dumb"`` or ``"unknown"`` will disable color/style and some features that require moving the cursor, such as progress bars.
-
-If the environment variable ``NO_COLOR`` is set, Rich will disable all color in the output.
-
-
 Printing
 --------
 
@@ -220,3 +178,52 @@ The Console class can export anything written to it as either text or html. To e
 After you have written content, you can call :meth:`~rich.console.Console.export_text` or :meth:`~rich.console.Console.export_html` to get the console output as a string. You can also call :meth:`~rich.console.Console.save_text` or :meth:`~rich.console.Console.save_html` to write the contents directly to disk.
 
 For examples of the html output generated by Rich Console, see :ref:`appendix-colors`.
+
+
+File output
+-----------
+
+The Console object will write to standard output (i.e. the terminal). You can also tell the Console object to write to another file by setting the ``file`` argument on the constructor -- which should be a file-like object opened for writing text. One use of this capability is to create a Console for writing to standard error by setting file to ``sys.stderr``. Here's an example::
+
+    import sys
+    from rich.console import Console
+    error_console = Console(file=sys.stderr)
+    error_console.print("[bold red]This is an error!")
+
+
+Capturing output
+----------------
+
+There may be situations where you want to *capture* the output from a Console rather than writing it directly to the terminal. You can do this with the :meth:`~rich.console.Console.capture` method which returns a context manager. On exit from this context manager, call :meth:`~rich.console.Capture.get` to return the string that would have been written to the terminal. Here's an example::
+
+    from rich.console import Console
+    console = Console()
+    with console.capture() as capture:
+        console.print("[bold red]Hello[/] World")
+    str_output = capture.get()
+
+An alternative way of capturing output is to set the Console file to a :py:class:`io.StringIO`. This is the recommended method if you are testing console output in unit tests. Here's an example::
+
+    from io import StringIO
+    from rich.console import Console
+    console = Console(file=StringIO())
+    console.print("[bold red]Hello[/] World")
+    str_output = console.file.getvalue()
+
+
+Terminal detection
+------------------
+
+If Rich detects that it is not writing to a terminal it will strip control codes from the output. If you want to write control codes to a regular file then set ``force_terminal=True`` on the constructor.
+
+Letting Rich auto-detect terminals is useful as it will write plain text when you pipe output to a file or other application.
+
+
+Environment variables
+---------------------
+
+Rich respects some standard environment variables.
+
+Setting the environment variable ``TERM`` to ``"dumb"`` or ``"unknown"`` will disable color/style and some features that require moving the cursor, such as progress bars.
+
+If the environment variable ``NO_COLOR`` is set, Rich will disable all color in the output.

rich/color.py

@@ -508,7 +508,7 @@ def blend_rgb(
 if __name__ == "__main__":  # pragma: no cover
 
     from .console import Console
-    from .table import Column, Table
+    from .table import Table
     from .text import Text
     from . import box
 

rich/color_triplet.py

@@ -22,7 +22,7 @@ class ColorTriplet(NamedTuple):
         """The color in RGB format.
 
         Returns:
-            str: An rgb color, e.g. ``"rgb(100, 23, 255)"``.
+            str: An rgb color, e.g. ``"rgb(100,23,255)"``.
         """
         red, green, blue = self
         return f"rgb({red},{green},{blue})"

rich/console.py

@@ -150,6 +150,38 @@ RenderResult = Iterable[Union[RenderableType, Segment]]
 _null_highlighter = NullHighlighter()
 
 
+class CaptureError(Exception):
+    """An error in the Capture context manager."""
+
+
+class Capture:
+    """Context manager to capture the result of printing to the console.
+    See :meth:`~rich.console.Console.capture` for how to use.
+
+    Args:
+        console (Console): A console instance to capture output.
+    """
+
+    def __init__(self, console: "Console") -> None:
+        self._console = console
+        self._result: Optional[str] = None
+
+    def __enter__(self) -> "Capture":
+        self._console.begin_capture()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
+        self._result = self._console.end_capture()
+
+    def get(self) -> str:
+        """Get the result of the capture."""
+        if self._result is None:
+            raise CaptureError(
+                "Capture result is not available until context manager exits."
+            )
+        return self._result
+
+
 class RenderGroup:
     """Takes a group of renderables and returns a renderable object that renders the group.
 
@@ -451,6 +483,21 @@ class Console:
         """Exit buffer context."""
         self._exit_buffer()
 
+    def begin_capture(self) -> None:
+        """Begin capturing console output. Call :meth:`end_capture` to exit capture mode and return output."""
+        self._enter_buffer()
+
+    def end_capture(self) -> str:
+        """End capture mode and return captured string.
+
+        Returns:
+            str: Console output.
+        """
+        render_result = self._render_buffer()
+        del self._buffer[:]
+        self._exit_buffer()
+        return render_result
+
     @property
     def color_system(self) -> Optional[str]:
         """Get color system string.
@@ -541,6 +588,23 @@ class Console:
         width, _ = self.size
         return width
 
+    def capture(self) -> Capture:
+        """A context manager to *capture* the result of print() or log() in a string,
+        rather than writing it to the console.
+
+        Example:
+            >>> from rich.console import Console
+            >>> console = Console()
+            >>> with console.capture() as capture:
+            ...     console.print("[bold magenta]Hello World[/])
+            >>> print(capture.result)
+
+        Returns:
+            Capture: Context manager which will contain the attribute `result` on exit.
+        """
+        capture = Capture(self)
+        return capture
+
     def line(self, count: int = 1) -> None:
         """Write new line(s).
 
@@ -1085,13 +1149,19 @@ class Console:
         Returns:
             str: Text read from stdin.
         """
+        prompt_str = ""
         if prompt:
-            self.print(prompt, markup=markup, emoji=emoji, end="")
-        result = (
-            getpass("", stream=stream)
-            if password
-            else (stream.readline() if stream else input())
-        )
+            with self.capture() as capture:
+                self.print(prompt, markup=markup, emoji=emoji, end="")
+            prompt_str = capture.get()
+        if password:
+            result = getpass(prompt_str, stream=stream)
+        else:
+            if stream:
+                self.file.write(prompt_str)
+                result = stream.readline()
+            else:
+                result = input(prompt_str)
         return result
 
     def export_text(self, *, clear: bool = True, styles: bool = False) -> str:

rich/containers.py

@@ -156,7 +156,7 @@ class Lines:
                     tokens.append(word)
                     if index < len(spaces):
                         if next_word is None:
-                            space_style = Style()
+                            space_style = Style.empty()
                         else:
                             style = word.get_style_at_offset(console, -1)
                             next_style = next_word.get_style_at_offset(console, 0)

rich/markup.py

@@ -31,6 +31,7 @@ class Tag(NamedTuple):
 
     @property
     def markup(self) -> str:
+        """Get the string representation of this tag."""
         return (
             f"[{self.name}]"
             if self.parameters is None

rich/style.py

@@ -162,6 +162,21 @@ class Style:
             )
         )
 
+    @classmethod
+    def empty(cls) -> "Style":
+        """Create an 'empty' style, equivalent to Style(), but more performant."""
+        style = cls.__new__(Style)
+        style._ansi = None
+        style._style_definition = None
+        style._color = None
+        style._bgcolor = None
+        style._attributes = 0
+        style._set_attributes = 0
+        style._link = None
+        style._link_id = None
+        style._hash = hash((None, None, 0, 0, None))
+        return style
+
     bold = _Bit(0)
     dim = _Bit(1)
     italic = _Bit(2)
@@ -532,7 +547,6 @@ class Style:
             rendered = (
                 f"\x1b]8;id={self._link_id};{self._link}\x1b\\{rendered}\x1b]8;;\x1b\\"
             )
-
         return rendered
 
     def test(self, text: Optional[str] = None) -> None:

rich/syntax.py

@@ -136,7 +136,7 @@ class PygmentsSyntaxTheme(SyntaxTheme):
             try:
                 pygments_style = self._pygments_style_class.style_for_token(token_type)
             except KeyError:
-                style = Style()
+                style = Style.empty()
             else:
                 color = pygments_style["color"]
                 bgcolor = pygments_style["bgcolor"]
@@ -159,8 +159,8 @@ class ANSISyntaxTheme(SyntaxTheme):
 
     def __init__(self, style_map: Dict[TokenType, Style]) -> None:
         self.style_map = style_map
-        self._missing_style = Style()
-        self._background_style = Style()
+        self._missing_style = Style.empty()
+        self._background_style = Style.empty()
         self._style_cache: Dict[TokenType, Style] = {}
 
     def get_style_for_token(self, token_type: TokenType) -> Style:
@@ -397,7 +397,7 @@ class Syntax(JupyterMixin):
         """Get background, number, and highlight styles for line numbers."""
         background_style = self._get_base_style()
         if background_style.transaprent_background:
-            return Style(), Style(dim=True), Style()
+            return Style.empty(), Style(dim=True), Style.empty()
         if console.color_system in ("256", "truecolor"):
             number_style = Style.chain(
                 background_style,

rich/table.py

@@ -229,7 +229,7 @@ class Table(JupyterMixin):
         """Get the current row style."""
         if self.row_styles:
             return self.row_styles[index % len(self.row_styles)]
-        return Style()
+        return Style.empty()
 
     def __rich_measure__(self, console: "Console", max_width: int) -> Measurement:
         if self.width is not None:
@@ -643,7 +643,7 @@ class Table(JupyterMixin):
             max_height = 1
             cells: List[List[List[Segment]]] = []
             if header_row or footer_row:
-                row_style = Style()
+                row_style = Style.empty()
             else:
                 row_style = get_style(
                     get_row_style(index - 1 if show_header else index)

rich/text.py

@@ -522,7 +522,7 @@ class Text(JupyterMixin):
         """
 
         text = self.plain
-        null_style = Style()
+        null_style = Style.empty()
         enumerated_spans = list(enumerate(self._spans, 1))
         get_style = partial(console.get_style, default=null_style)
         style_map = {index: get_style(span.style) for index, span in enumerated_spans}

tests/test_console.py

@@ -5,13 +5,11 @@ import tempfile
 
 import pytest
 
-from rich.color import Color, ColorSystem
-from rich.console import Console, ConsoleOptions
+from rich.color import ColorSystem
+from rich.console import CaptureError, Console, ConsoleOptions
 from rich import errors
 from rich.panel import Panel
-from rich.segment import Segment
 from rich.style import Style
-from rich.theme import Theme
 
 
 def test_dumb_terminal():
@@ -160,14 +158,41 @@ def test_control():
     assert console.file.getvalue() == "FOOBAR\n"
 
 
+def test_capture():
+    console = Console()
+    with console.capture() as capture:
+        with pytest.raises(CaptureError):
+            capture.get()
+        console.print("Hello")
+    assert capture.get() == "Hello\n"
+
+
 def test_input(monkeypatch, capsys):
-    monkeypatch.setattr("builtins.input", lambda: "bar")
+    def fake_input(prompt):
+        console.file.write(prompt)
+        return "bar"
+
+    monkeypatch.setattr("builtins.input", fake_input)
     console = Console()
     user_input = console.input(prompt="foo:")
     assert capsys.readouterr().out == "foo:"
     assert user_input == "bar"
 
 
+def test_input_password(monkeypatch, capsys):
+    def fake_input(prompt, stream=None):
+        console.file.write(prompt)
+        return "bar"
+
+    import rich.console
+
+    monkeypatch.setattr(rich.console, "getpass", fake_input)
+    console = Console()
+    user_input = console.input(prompt="foo:", password=True)
+    assert capsys.readouterr().out == "foo:"
+    assert user_input == "bar"
+
+
 def test_justify_none():
     console = Console(file=io.StringIO(), force_terminal=True, width=20)
     console.print("FOO", justify=None)

tests/test_style.py

@@ -71,6 +71,10 @@ def test_hash():
     assert isinstance(hash(Style()), int)
 
 
+def test_empty():
+    assert Style.empty() == Style()
+
+
 def test_bool():
     assert bool(Style()) is False
     assert bool(Style(bold=True)) is True