From ebfab2ddcb45a16b3a6f0cbfabb66cd5a6b576be Mon Sep 17 00:00:00 2001
From: Will McGugan <willmcgugan@gmail.com>
Date: Fri, 2 Apr 2021 11:29:48 +0100
Subject: [PATCH] better Jupyter support

---
 rich/console.py |  2 +-
 rich/jupyter.py | 37 ++++++++++++++++++++-----------------
 rich/layout.py  |  2 --
 rich/pretty.py  |  8 +++++---
 4 files changed, 26 insertions(+), 23 deletions(-)

diff --git a/rich/console.py b/rich/console.py
index c01aee3d..71fa0809 100644
--- a/rich/console.py
+++ b/rich/console.py
@@ -1681,7 +1681,7 @@ class Console:
                 if self.is_jupyter:  # pragma: no cover
                     from .jupyter import display
 
-                    display(self._buffer)
+                    display(self._buffer, self._render_buffer(self._buffer[:]))
                     del self._buffer[:]
                 else:
                     text = self._render_buffer(self._buffer[:])
diff --git a/rich/jupyter.py b/rich/jupyter.py
index 9101ce14..c4d66569 100644
--- a/rich/jupyter.py
+++ b/rich/jupyter.py
@@ -1,11 +1,9 @@
-from typing import Iterable, List, TYPE_CHECKING
+from typing import Iterable, List
 
 from . import get_console
 from .segment import Segment
 from .terminal_theme import DEFAULT_TERMINAL_THEME
 
-if TYPE_CHECKING:
-    from .console import RenderableType
 
 JUPYTER_HTML_FORMAT = """\
 <pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace">{code}</pre>
@@ -15,28 +13,33 @@ JUPYTER_HTML_FORMAT = """\
 class JupyterRenderable:
     """A shim to write html to Jupyter notebook."""
 
-    def __init__(self, html: str) -> None:
+    def __init__(self, html: str, text: str) -> None:
         self.html = html
+        self.text = text
 
-    @classmethod
-    def render(cls, rich_renderable: "RenderableType") -> str:
-        console = get_console()
-        segments = console.render(rich_renderable, console.options)
-        html = _render_segments(segments)
-        return html
-
-    def _repr_html_(self) -> str:
-        return self.html
+    def _repr_mimebundle_(self, include, exclude, **kwargs):
+        data = {"text/plain": self.text, "text/html": self.html}
+        if include:
+            data = {k: v for (k, v) in data.items() if k in include}
+        if exclude:
+            data = {k: v for (k, v) in data.items() if k not in exclude}
+        return data
 
 
 class JupyterMixin:
     """Add to an Rich renderable to make it render in Jupyter notebook."""
 
-    def _repr_html_(self) -> str:
+    def _repr_mimebundle_(self, include, exclude, **kwargs):
         console = get_console()
         segments = list(console.render(self, console.options))  # type: ignore
         html = _render_segments(segments)
-        return html
+        text = console._render_buffer(segments)
+        data = {"text/plain": text, "text/html": html}
+        if include:
+            data = {k: v for (k, v) in data.items() if k in include}
+        if exclude:
+            data = {k: v for (k, v) in data.items() if k not in exclude}
+        return data
 
 
 def _render_segments(segments: Iterable[Segment]) -> str:
@@ -64,12 +67,12 @@ def _render_segments(segments: Iterable[Segment]) -> str:
     return html
 
 
-def display(segments: Iterable[Segment]) -> None:
+def display(segments: Iterable[Segment], text: str) -> None:
     """Render segments to Jupyter."""
     from IPython.display import display as ipython_display
 
     html = _render_segments(segments)
-    jupyter_renderable = JupyterRenderable(html)
+    jupyter_renderable = JupyterRenderable(html, text)
     ipython_display(jupyter_renderable)
 
 
diff --git a/rich/layout.py b/rich/layout.py
index c8aab23b..21057828 100644
--- a/rich/layout.py
+++ b/rich/layout.py
@@ -14,8 +14,6 @@ from typing import (
     Union,
 )
 
-from typing_extensions import Literal
-
 from ._ratio import ratio_resolve
 from .align import Align
 from .console import Console, ConsoleOptions, RenderableType, RenderResult
diff --git a/rich/pretty.py b/rich/pretty.py
index 7c9f8e1a..7903b2cc 100644
--- a/rich/pretty.py
+++ b/rich/pretty.py
@@ -4,7 +4,6 @@ import sys
 from array import array
 from collections import Counter, defaultdict, deque
 from dataclasses import dataclass, fields, is_dataclass
-import inspect
 from itertools import islice
 from typing import (
     TYPE_CHECKING,
@@ -27,7 +26,7 @@ from ._pick import pick_bool
 from .abc import RichRenderable
 from .cells import cell_len
 from .highlighter import ReprHighlighter
-from .jupyter import JupyterRenderable
+from .jupyter import JupyterMixin, JupyterRenderable
 from .measure import Measurement
 from .text import Text
 
@@ -99,6 +98,9 @@ def install(
         if console.is_jupyter and any(attr.startswith("_repr_") for attr in dir(value)):
             return
 
+        if hasattr(value, "_repr_mimebundle_"):
+            return
+
         # certain renderables should start on a new line
         if isinstance(value, ConsoleRenderable):
             console.line()
@@ -130,7 +132,7 @@ def install(
         sys.displayhook = display_hook
 
 
-class Pretty:
+class Pretty(JupyterMixin):
     """A rich renderable that pretty prints an object.
 
     Args: