pw_console: ptpython repl, execution and display

Enables a ptpython based repl in the ReplPane. User repl code is
executed in it's own thread with stdout and stderr patched to
capture output.

Change-Id: I005481f1ecd6805ce9a74fbf57ff0f2317b9d2aa
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/48960
Pigweed-Auto-Submit: Anthony DiGirolamo <tonymd@google.com>
Commit-Queue: Auto-Submit <auto-submit@pigweed.google.com.iam.gserviceaccount.com>
Reviewed-by: Joe Ethier <jethier@google.com>
diff --git a/pw_console/docs.rst b/pw_console/docs.rst
index 0f41a3a..51f1bdd 100644
--- a/pw_console/docs.rst
+++ b/pw_console/docs.rst
@@ -5,22 +5,34 @@
 ----------
 
 The Pigweed Console provides a Python repl (read eval print loop) using
-`ptpython <https://github.com/prompt-toolkit/ptpython>`_ and a log message
-viewer in a single-window terminal based interface. It is designed to be a
-replacement for
-`IPython's embed() <https://ipython.readthedocs.io/en/stable/interactive/reference.html#embedding>`_
-function.
+`ptpython`_ and a log message viewer in a single-window terminal based
+interface. It is designed to be a replacement for `IPython's embed()`_ function.
+
+.. warning::
+   The Pigweed Console is under heavy development. A user manual and usage
+   information will be documented as features near completion.
 
 ==========
-Motivation
+Goals
 ==========
 
-``pw_console`` is the complete solution for interacting with hardware devices
+``pw_console`` is a complete solution for interacting with hardware devices
 using :ref:`module-pw_rpc` over a :ref:`module-pw_hdlc` transport.
 
 The repl allows interactive RPC sending while the log viewer provides immediate
 feedback on device status.
 
+- Interactive Python repl and log viewer in a single terminal window.
+
+- Easily embeddable within a project's own custom console. This should allow
+  users to define their own transport layer.
+
+- Plugin framework to add custom status toolbars or window panes.
+
+- Log viewer with searching and filtering.
+
+- Daemon that provides a many-to-many mapping between consoles and devices.
+
 =====
 Usage
 =====
@@ -42,7 +54,18 @@
 Thread and Event Loop Design
 ----------------------------
 
-Here's a diagram showing how ``pw_console`` threads and asyncio tasks are organized.
+In `ptpython`_ and `IPython`_ all user repl code is run in the foreground. This
+allows interrupts like ``Ctrl-C`` and functions like ``print()`` and
+``time.sleep()`` to work as expected. Pigweed's Console doesn't use this
+approach as it would hide or freeze the `prompt_toolkit`_ user interface while
+running repl code.
+
+To get around this issue all user repl code is run in a dedicated thread with
+stdout and stderr patched to capture output. This lets the user interface stay
+responsive and new log messages to continue to be displayed.
+
+Here's a diagram showing how ``pw_console`` threads and `asyncio`_ tasks are
+organized.
 
 .. mermaid::
 
@@ -96,3 +119,9 @@
        repl-.->|Run Code| replThread
        pluginToolbar-.->|Register Plugin| pluginThread
        pluginPane-.->|Register Plugin| pluginThread2
+
+.. _IPython's embed(): https://ipython.readthedocs.io/en/stable/interactive/reference.html#embedding
+.. _IPython: https://ipython.readthedocs.io/
+.. _asyncio: https://docs.python.org/3/library/asyncio.html
+.. _prompt_toolkit: https://python-prompt-toolkit.readthedocs.io/
+.. _ptpython: https://github.com/prompt-toolkit/ptpython/
diff --git a/pw_console/py/BUILD.gn b/pw_console/py/BUILD.gn
index 98ec217..10865c2 100644
--- a/pw_console/py/BUILD.gn
+++ b/pw_console/py/BUILD.gn
@@ -23,14 +23,17 @@
     "pw_console/__main__.py",
     "pw_console/console_app.py",
     "pw_console/help_window.py",
+    "pw_console/helpers.py",
     "pw_console/key_bindings.py",
     "pw_console/log_pane.py",
+    "pw_console/pw_ptpython_repl.py",
     "pw_console/repl_pane.py",
     "pw_console/style.py",
   ]
   tests = [
     "console_app_test.py",
     "help_window_test.py",
+    "repl_pane_test.py",
   ]
   python_deps = [
     "$dir_pw_cli/py",
diff --git a/pw_console/py/help_window_test.py b/pw_console/py/help_window_test.py
index 55feb3d..d7352b6 100644
--- a/pw_console/py/help_window_test.py
+++ b/pw_console/py/help_window_test.py
@@ -23,7 +23,7 @@
 
 
 class TestHelpWindow(unittest.TestCase):
-    """Tests for ConsoleApp."""
+    """Tests for HelpWindow text and keybind lists."""
     def setUp(self):
         self.maxDiff = None  # pylint: disable=invalid-name
 
diff --git a/pw_console/py/pw_console/console_app.py b/pw_console/py/pw_console/console_app.py
index ae50e93..d9444f5 100644
--- a/pw_console/py/pw_console/console_app.py
+++ b/pw_console/py/pw_console/console_app.py
@@ -38,10 +38,12 @@
     MenuItem,
 )
 from prompt_toolkit.key_binding import merge_key_bindings
+from ptpython.key_bindings import load_python_bindings  # type: ignore
 
 from pw_console.help_window import HelpWindow
 from pw_console.key_bindings import create_key_bindings
 from pw_console.log_pane import LogPane
+from pw_console.pw_ptpython_repl import PwPtPythonRepl
 from pw_console.repl_pane import ReplPane
 from pw_console.style import pw_console_styles
 
@@ -94,9 +96,19 @@
         self.show_help_window = False
         self.vertical_split = False
 
-        # Create one log pane and the repl pane.
+        # Create one log pane.
         self.log_pane = LogPane(application=self)
-        self.repl_pane = ReplPane(application=self)
+
+        # Create a ptpython repl instance.
+        self.pw_ptpython_repl = PwPtPythonRepl(
+            get_globals=lambda: global_vars,
+            get_locals=lambda: local_vars,
+        )
+
+        self.repl_pane = ReplPane(
+            application=self,
+            python_repl=self.pw_ptpython_repl,
+        )
 
         # List of enabled panes.
         self.active_panes = [
@@ -162,14 +174,14 @@
             layout=self.layout,
             after_render=self.run_after_render_hooks,
             key_bindings=merge_key_bindings([
-                # TODO: pull key bindings from ptpython
-                # load_python_bindings(self.pw_ptpython_repl),
+                # Pull key bindings from ptpython
+                load_python_bindings(self.pw_ptpython_repl),
                 self.key_bindings,
             ]),
             style=DynamicStyle(lambda: merge_styles([
                 pw_console_styles,
-                # TODO: Include ptpython styles
-                # self.pw_ptpython_repl._current_style
+                # Include ptpython styles
+                self.pw_ptpython_repl._current_style,  # pylint: disable=protected-access
             ])),
             enable_page_navigation_bindings=True,
             full_screen=True,
diff --git a/pw_console/py/pw_console/helpers.py b/pw_console/py/pw_console/helpers.py
new file mode 100644
index 0000000..484009b
--- /dev/null
+++ b/pw_console/py/pw_console/helpers.py
@@ -0,0 +1,19 @@
+# Copyright 2021 The Pigweed Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); you may not
+# use this file except in compliance with the License. You may obtain a copy of
+# the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations under
+# the License.
+"""Helper functions."""
+
+
+def remove_formatting(formatted_text):
+    """Throw away style info from formatted text tuples."""
+    return ''.join([formatted_tuple[1] for formatted_tuple in formatted_text])  # pylint: disable=not-an-iterable
diff --git a/pw_console/py/pw_console/key_bindings.py b/pw_console/py/pw_console/key_bindings.py
index 3579d17..2e527b4 100644
--- a/pw_console/py/pw_console/key_bindings.py
+++ b/pw_console/py/pw_console/key_bindings.py
@@ -56,6 +56,8 @@
     @bindings.add('c-q')
     def exit_(event):
         """Quit the console application."""
+        # TODO(tonymd): Cancel any existing user repl or plugin tasks before
+        # exiting.
         event.app.exit()
 
     @bindings.add('s-tab')
@@ -71,4 +73,17 @@
         """Move focus to the previous widget."""
         focus_previous(event)
 
+    # Bindings for when the ReplPane input field is in focus.
+    @bindings.add('c-c', filter=has_focus(console_app.pw_ptpython_repl))
+    def handle_ctrl_c(event):
+        """Reset the python repl on Ctrl-c"""
+        console_app.repl_pane.ctrl_c()
+
+    @bindings.add('c-d', filter=has_focus(console_app.pw_ptpython_repl))
+    def handle_ctrl_d(event):
+        """Do nothing on ctrl-d."""
+        # TODO(tonymd): Allow ctrl-d to quit the whole app with confirmation
+        # like ipython.
+        pass
+
     return bindings
diff --git a/pw_console/py/pw_console/pw_ptpython_repl.py b/pw_console/py/pw_console/pw_ptpython_repl.py
new file mode 100644
index 0000000..bcc9fc1
--- /dev/null
+++ b/pw_console/py/pw_console/pw_ptpython_repl.py
@@ -0,0 +1,185 @@
+# Copyright 2021 The Pigweed Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); you may not
+# use this file except in compliance with the License. You may obtain a copy of
+# the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations under
+# the License.
+"""PwPtPythonPane class."""
+
+import asyncio
+import logging
+import io
+import sys
+from functools import partial
+from pathlib import Path
+
+from prompt_toolkit.buffer import Buffer
+import ptpython.repl  # type: ignore
+
+from pw_console.helpers import remove_formatting
+
+_LOG = logging.getLogger(__package__)
+
+
+class PwPtPythonRepl(ptpython.repl.PythonRepl):
+    """A ptpython repl class with changes to code execution and output related
+    methods."""
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args,
+                         create_app=False,
+                         history_filename=(Path.home() /
+                                           '.pw_console_history').as_posix(),
+                         color_depth='256 colors',
+                         _input_buffer_height=8,
+                         **kwargs)
+        # Change some ptpython.repl defaults.
+        self.use_code_colorscheme('zenburn')
+        self.show_status_bar = False
+        self.show_exit_confirmation = False
+        self.complete_private_attributes = False
+
+        # Additional state variables.
+        self.repl_pane = None
+        self._last_result = None
+
+    def __pt_container__(self):
+        """Return the prompt_toolkit root container for class."""
+        return self.ptpython_layout.root_container
+
+    def set_repl_pane(self, repl_pane):
+        """Update the parent pw_console.ReplPane reference."""
+        self.repl_pane = repl_pane
+
+    def _save_result(self, formatted_text):
+        """Save the last repl execution result."""
+        unformatted_result = remove_formatting(formatted_text)
+        self._last_result = unformatted_result
+
+    def clear_last_result(self):
+        """Erase the last repl execution result."""
+        self._last_result = None
+
+    def _update_output_buffer(self):
+        self.repl_pane.update_output_buffer()
+
+    def show_result(self, result):
+        """Format and save output results.
+
+        This function is called from the _run_user_code() function which is
+        always run from the user code thread, within
+        .run_and_show_expression_async().
+        """
+        formatted_result = self._format_result_output(result)
+        self._save_result(formatted_result)
+
+    def _handle_exception(self, e: BaseException) -> None:
+        """Format and save output results.
+
+        This function is called from the _run_user_code() function which is
+        always run from the user code thread, within
+        .run_and_show_expression_async().
+        """
+        formatted_result = self._format_exception_output(e)
+        self._save_result(formatted_result.__pt_formatted_text__())
+
+    def user_code_complete_callback(self, input_text, future):
+        """Callback to run after user repl code is finished."""
+        # If there was an exception it will be saved in self._last_result
+        result = self._last_result
+        # _last_result consumed, erase for the next run.
+        self.clear_last_result()
+
+        stdout_contents = None
+        stderr_contents = None
+        if future.result():
+            future_result = future.result()
+            stdout_contents = future_result['stdout']
+            stderr_contents = future_result['stderr']
+            result_value = future_result['result']
+
+            if result_value is not None:
+                formatted_result = self._format_result_output(result_value)
+                result = remove_formatting(formatted_result)
+
+        # Job is finished, append the last result.
+        self.repl_pane.append_result_to_executed_code(input_text, future,
+                                                      result, stdout_contents,
+                                                      stderr_contents)
+
+        # Rebuild output buffer.
+        self._update_output_buffer()
+
+        # Trigger a prompt_toolkit application redraw.
+        self.repl_pane.application.application.invalidate()
+
+    async def _run_user_code(self, text):
+        """Run user code and capture stdout+err.
+
+        This fuction should be run in a separate thread from the main
+        prompt_toolkit application."""
+        # NOTE: This function runs in a separate thread using the asyncio event
+        # loop defined by self.repl_pane.application.user_code_loop. Patching
+        # stdout here will not effect the stdout used by prompt_toolkit and the
+        # main user interface.
+
+        # Patch stdout and stderr to capture repl print() statements.
+        original_stdout = sys.stdout
+        original_stderr = sys.stderr
+
+        temp_out = io.StringIO()
+        temp_err = io.StringIO()
+
+        sys.stdout = temp_out
+        sys.stderr = temp_err
+
+        # Run user repl code
+        try:
+            result = await self.run_and_show_expression_async(text)
+        finally:
+            # Always restore original stdout and stderr
+            sys.stdout = original_stdout
+            sys.stderr = original_stderr
+
+        # Save the captured output
+        stdout_contents = temp_out.getvalue()
+        stderr_contents = temp_err.getvalue()
+
+        return {
+            'stdout': stdout_contents,
+            'stderr': stderr_contents,
+            'result': result
+        }
+
+    def _accept_handler(self, buff: Buffer) -> bool:
+        """Function executed when pressing enter in the ptpython.repl.PythonRepl
+        input buffer."""
+        # Do nothing if no text is entered.
+        if len(buff.text) == 0:
+            return False
+
+        # Execute the repl code in the the separate user_code thread loop.
+        future = asyncio.run_coroutine_threadsafe(
+            # This function will be executed in a separate thread.
+            self._run_user_code(buff.text),
+            # Using this asyncio event loop.
+            self.repl_pane.application.user_code_loop)
+        # Run user_code_complete_callback() when done.
+        done_callback = partial(self.user_code_complete_callback, buff.text)
+        future.add_done_callback(done_callback)
+
+        # Save the input text and future object.
+        self.repl_pane.append_executed_code(buff.text, future)
+
+        # Rebuild the parent ReplPane output buffer.
+        self._update_output_buffer()
+
+        # TODO: Return True if exception is found?
+        # Don't keep input for now. Return True to keep input text.
+        return False
diff --git a/pw_console/py/pw_console/repl_pane.py b/pw_console/py/pw_console/repl_pane.py
index 5d301ca..86ac717 100644
--- a/pw_console/py/pw_console/repl_pane.py
+++ b/pw_console/py/pw_console/repl_pane.py
@@ -13,8 +13,11 @@
 # the License.
 """ReplPane class."""
 
+import concurrent
 import logging
+from dataclasses import dataclass
 from functools import partial
+from pathlib import Path
 from typing import (
     Any,
     Callable,
@@ -23,10 +26,12 @@
     Optional,
 )
 
+from jinja2 import Template
 from prompt_toolkit.filters import (
     Condition,
     has_focus,
 )
+from prompt_toolkit.document import Document
 from prompt_toolkit.mouse_events import MouseEvent, MouseEventType
 from prompt_toolkit.layout.dimension import AnyDimension
 from prompt_toolkit.widgets import TextArea
@@ -44,11 +49,18 @@
 from prompt_toolkit.lexers import PygmentsLexer  # type: ignore
 from pygments.lexers.python import PythonLexer  # type: ignore
 
+from pw_console.pw_ptpython_repl import PwPtPythonRepl
+
 _LOG = logging.getLogger(__package__)
 
 _Namespace = Dict[str, Any]
 _GetNamespace = Callable[[], _Namespace]
 
+_OUTPUT_TEMPLATE_PATH = (Path(__file__).parent / 'templates' /
+                         'repl_output.jinja')
+with _OUTPUT_TEMPLATE_PATH.open() as tmpl:
+    OUTPUT_TEMPLATE = tmpl.read()
+
 
 def mouse_focus_handler(repl_pane, mouse_event: MouseEvent):
     """Focus the repl_pane on click."""
@@ -89,9 +101,9 @@
         """Return toolbar text showing if the ReplPane is in focus or not."""
         focused_text = (
             # Style
-            "",
+            '',
             # Text
-            " [FOCUSED] ",
+            ' [FOCUSED] ',
             # Mouse handler
             partial(mouse_focus_handler, repl_pane),
         )
@@ -171,6 +183,20 @@
             filter=Condition(lambda: repl_pane.show_bottom_toolbar))
 
 
+@dataclass
+class UserCodeExecution:
+    """Class to hold a single user repl execution event."""
+    input: str
+    future: concurrent.futures.Future
+    output: str
+    stdout: str
+    stderr: str
+
+    @property
+    def is_running(self):
+        return not self.future.done()
+
+
 class ReplPane:
     """Pane for reading Python input."""
 
@@ -178,8 +204,7 @@
     def __init__(
             self,
             application: Any,
-            # TODO: Include ptpython repl.
-            # python_repl: PwPtPythonRepl,
+            python_repl: PwPtPythonRepl,
             # TODO: Make the height of input+output windows match the log pane
             # height. (Using minimum output height of 5 for now).
             output_height: Optional[AnyDimension] = Dimension(preferred=5),
@@ -197,13 +222,10 @@
         self.show_top_toolbar = True
         self.show_bottom_toolbar = True
 
-        # TODO: Include ptpython repl.
-        self.pw_ptpython_repl = Window(content=FormattedTextControl(
-            [('', '>>> Repl input buffer')], focusable=True))
-        self.last_error_output = ""
+        self.pw_ptpython_repl = python_repl
+        self.pw_ptpython_repl.set_repl_pane(self)
 
         self.output_field = TextArea(
-            text='Repl output buffer',
             style='class:output-field',
             height=output_height,
             # text=help_text,
@@ -265,3 +287,75 @@
 
     def after_render_hook(self):
         """Run tasks after the last UI render."""
+
+    def ctrl_c(self):
+        """Ctrl-C keybinding behavior."""
+        # If there is text in the input buffer, clear it.
+        if self.pw_ptpython_repl.default_buffer.text:
+            self.clear_input_buffer()
+        else:
+            self.interrupt_last_code_execution()
+
+    def clear_input_buffer(self):
+        self.pw_ptpython_repl.default_buffer.reset()
+
+    def interrupt_last_code_execution(self):
+        code = self._get_currently_running_code()
+        if code:
+            code.future.cancel()
+            code.output = 'Canceled'
+        self.pw_ptpython_repl.clear_last_result()
+        self.update_output_buffer()
+
+    def _get_currently_running_code(self):
+        for code in self.executed_code:
+            if not code.future.done():
+                return code
+        return None
+
+    def _get_executed_code(self, future):
+        for code in self.executed_code:
+            if code.future == future:
+                return code
+        return None
+
+    def _log_executed_code(self, code, prefix=''):
+        text = self.get_output_buffer_text([code], show_index=False)
+        _LOG.info('[PYTHON] %s\n%s', prefix, text)
+
+    def append_executed_code(self, text, future):
+        user_code = UserCodeExecution(input=text,
+                                      future=future,
+                                      output=None,
+                                      stdout=None,
+                                      stderr=None)
+        self.executed_code.append(user_code)
+        self._log_executed_code(user_code, prefix='START')
+
+    def append_result_to_executed_code(self,
+                                       _input_text,
+                                       future,
+                                       result_text,
+                                       stdout_text='',
+                                       stderr_text=''):
+        code = self._get_executed_code(future)
+        if code:
+            code.output = result_text
+            code.stdout = stdout_text
+            code.stderr = stderr_text
+        self._log_executed_code(code, prefix='FINISH')
+        self.update_output_buffer()
+
+    def get_output_buffer_text(self, code_items=None, show_index=True):
+        executed_code = code_items or self.executed_code
+        template = Template(OUTPUT_TEMPLATE,
+                            trim_blocks=True,
+                            lstrip_blocks=True)
+        return template.render(code_items=executed_code,
+                               show_index=show_index).strip()
+
+    def update_output_buffer(self):
+        text = self.get_output_buffer_text()
+        self.output_field.buffer.document = Document(text=text,
+                                                     cursor_position=len(text))
+        self.application.redraw_ui()
diff --git a/pw_console/py/pw_console/templates/repl_output.jinja b/pw_console/py/pw_console/templates/repl_output.jinja
new file mode 100644
index 0000000..1ccf272
--- /dev/null
+++ b/pw_console/py/pw_console/templates/repl_output.jinja
@@ -0,0 +1,34 @@
+{#
+Copyright 2021 The Pigweed Authors
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not
+use this file except in compliance with the License. You may obtain a copy of
+the License at
+
+    https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+License for the specific language governing permissions and limitations under
+the License.
+#}
+{% for code in code_items %}
+{% set index = loop.index if show_index else '' %}
+{% set prompt_width = 7 + index|string|length %}
+In [{{index}}]: {{ code.input|indent(width=prompt_width) }}
+{% if code.is_running %}
+Running...
+{% else %}
+{% if code.stdout -%}
+  {{ code.stdout }}
+{%- endif %}
+{% if code.stderr -%}
+  {{ code.stderr }}
+{%- endif %}
+{% if code.output %}
+Out[{{index}}]: {{ code.output|indent(width=prompt_width) }}
+{% endif %}
+{% endif %}
+
+{% endfor -%}
diff --git a/pw_console/py/repl_pane_test.py b/pw_console/py/repl_pane_test.py
new file mode 100644
index 0000000..1a82925
--- /dev/null
+++ b/pw_console/py/repl_pane_test.py
@@ -0,0 +1,144 @@
+# Copyright 2021 The Pigweed Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License"); you may not
+# use this file except in compliance with the License. You may obtain a copy of
+# the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations under
+# the License.
+"""Tests for pw_console.console_app"""
+
+import asyncio
+import threading
+import builtins
+import unittest
+from inspect import cleandoc
+from unittest.mock import Mock, MagicMock
+
+from prompt_toolkit.application import create_app_session
+# inclusive-language: ignore
+from prompt_toolkit.output import DummyOutput as FakeOutput
+
+from pw_console.console_app import ConsoleApp
+from pw_console.repl_pane import ReplPane
+from pw_console.pw_ptpython_repl import PwPtPythonRepl
+
+
+class TestReplPane(unittest.TestCase):
+    """Tests for ReplPane."""
+    def test_repl_code_return_values(self) -> None:
+        """Test stdout, return values, and exceptions can be returned from
+        running user repl code."""
+        app = Mock()
+
+        global_vars = {
+            '__name__': '__main__',
+            '__package__': None,
+            '__doc__': None,
+            '__builtins__': builtins,
+        }
+
+        pw_ptpython_repl = PwPtPythonRepl(
+            get_globals=lambda: global_vars,
+            get_locals=lambda: global_vars,
+        )
+        repl_pane = ReplPane(
+            application=app,
+            python_repl=pw_ptpython_repl,
+        )
+        # Check pw_ptpython_repl has a reference to the parent repl_pane.
+        self.assertEqual(repl_pane, pw_ptpython_repl.repl_pane)
+
+        # Define a function, should return nothing.
+        code = cleandoc("""
+            def run():
+                print('The answer is ', end='')
+                return 1+1+4+16+20
+        """)
+        # pylint: disable=protected-access
+        result = asyncio.run(pw_ptpython_repl._run_user_code(code))
+        self.assertEqual(result, {'stdout': '', 'stderr': '', 'result': None})
+
+        # Check stdout and return value
+        result = asyncio.run(pw_ptpython_repl._run_user_code('run()'))
+        self.assertEqual(result, {
+            'stdout': 'The answer is ',
+            'stderr': '',
+            'result': 42
+        })
+
+        # Check for repl exception
+        result = asyncio.run(pw_ptpython_repl._run_user_code('return "blah"'))
+        self.assertIn("SyntaxError: 'return' outside function",
+                      pw_ptpython_repl._last_result)
+
+    def test_user_thread(self) -> None:
+        """Test user code thread."""
+        with create_app_session(output=FakeOutput()):
+            app = ConsoleApp()
+            app.start_user_code_thread()
+
+            pw_ptpython_repl = app.pw_ptpython_repl
+            repl_pane = app.repl_pane
+
+            pw_ptpython_repl.user_code_complete_callback = MagicMock(
+                wraps=pw_ptpython_repl.user_code_complete_callback)
+            user_code_done = threading.Event()
+
+            code = cleandoc("""
+                import time
+                def run():
+                    time.sleep(0.3)
+                    print('The answer is ', end='')
+                    return 1+1+4+16+20
+            """)
+
+            input_buffer = MagicMock(text=code)
+            # pylint: disable=protected-access
+            pw_ptpython_repl._accept_handler(input_buffer)
+
+            # Get last executed code object.
+            user_code1 = repl_pane.executed_code[-1]
+            # Wait for repl code to finish.
+            user_code1.future.add_done_callback(
+                lambda future: user_code_done.set())
+            user_code_done.wait(timeout=3)
+
+            pw_ptpython_repl.user_code_complete_callback.assert_called_once()
+            self.assertIsNotNone(user_code1)
+            self.assertTrue(user_code1.future.done())
+            self.assertEqual(user_code1.input, code)
+            self.assertEqual(user_code1.output, None)
+            # stdout / stderr may be '' or None
+            self.assertFalse(user_code1.stdout)
+            self.assertFalse(user_code1.stderr)
+
+            user_code_done.clear()
+            pw_ptpython_repl.user_code_complete_callback.reset_mock()
+
+            input_buffer = MagicMock(text='run()')
+            pw_ptpython_repl._accept_handler(input_buffer)
+
+            # Get last executed code object.
+            user_code2 = repl_pane.executed_code[-1]
+            # Wait for repl code to finish.
+            user_code2.future.add_done_callback(
+                lambda future: user_code_done.set())
+            user_code_done.wait(timeout=3)
+
+            pw_ptpython_repl.user_code_complete_callback.assert_called_once()
+            self.assertIsNotNone(user_code2)
+            self.assertTrue(user_code2.future.done())
+            self.assertEqual(user_code2.input, 'run()')
+            self.assertEqual(user_code2.output, '42')
+            self.assertEqual(user_code2.stdout, 'The answer is ')
+            self.assertFalse(user_code2.stderr)
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/pw_console/py/setup.py b/pw_console/py/setup.py
index 20d727d..c4b379d 100644
--- a/pw_console/py/setup.py
+++ b/pw_console/py/setup.py
@@ -26,6 +26,7 @@
         'pw_console': [
             'py.typed',
             'templates/keybind_list.jinja',
+            'templates/repl_output.jinja',
         ]
     },
     zip_safe=False,
@@ -39,7 +40,8 @@
         'ipython',
         'jinja2',
         'prompt_toolkit',
-        'ptpython',
+        # Required features are not yet in https://pypi.org/project/ptpython/
+        'ptpython @ git+https://github.com/prompt-toolkit/ptpython.git@b74af76',
         'pw_cli',
         'pw_tokenizer',
         'pygments',