Add prompt-toolkit theme integration for autocomplete and matching brackets

Author: dangooddd

README.md

@@ -24,7 +24,7 @@ pyrepl.setup({
   split_horizontal = false,
   split_ratio = 0.5,
   style = "default",
-  style_treesitter = true,
+  style_integration = true,
   image_max_history = 10,
   image_width_ratio = 0.5,
   image_height_ratio = 0.5,
@@ -150,11 +150,11 @@ python -m ipykernel install --user --name {kernel_name}
 
 ### Use a built-in Pygments style
 
-If you do not like the treesitter-based REPL colors, pick a built-in Pygments theme:
+If you do not want pyrepl.nvim to derive REPL colors from your Neovim theme, pick a built-in Pygments style instead:
 
 ```lua
 require("pyrepl").setup({
-  style_treesitter = false,
+  style_integration = false,
   style = "default", -- or another Pygments style, e.g. "gruvbox-dark"
 })
 ```

doc/pyrepl.txt

@@ -36,7 +36,7 @@ Minimal `vim.pack` setup with the default config and example keymaps:
     split_horizontal = false,
     split_ratio = 0.5,
     style = "default",
-    style_treesitter = true,
+    style_integration = true,
     image_max_history = 10,
     image_width_ratio = 0.5,
     image_height_ratio = 0.5,
@@ -206,7 +206,7 @@ Default configuration:
     split_horizontal = false,
     split_ratio = 0.5,
     style = "default",
-    style_treesitter = true,
+    style_integration = true,
     image_max_history = 10,
     image_width_ratio = 0.5,
     image_height_ratio = 0.5,
@@ -225,10 +225,10 @@ Options:
       Fraction of screen used by the REPL split. Default: 0.5.
   style (string)
       Pygments style name used by jupyter-console. Default: "default".
-      Used directly when `style_treesitter` is false.
-  style_treesitter (boolean)
-      Derive Pygments style overrides from current Neovim/Tree-sitter
-      highlight groups for REPL output. Default: true.
+      Used directly when `style_integration` is false.
+  style_integration (boolean)
+      Derive Pygments and prompt-toolkit style overrides from current
+      Neovim highlight groups for the REPL. Default: true.
       Set to false to use built-in Pygments styles from `style`.
   image_max_history (integer)
       Maximum number of images kept in history. Default: 10. Minimum: 2.
@@ -325,11 +325,11 @@ To use an arbitrary kernel in that case, you need to install it globally:
 
 Use a built-in Pygments style ~
 
-If you do not like the treesitter-based REPL colors, disable it and pick a
-built-in Pygments theme:
+If you do not want pyrepl.nvim to derive REPL colors from your Neovim
+highlights, disable the integration and pick a built-in Pygments theme:
 >
   require("pyrepl").setup({
-    style_treesitter = false,
+    style_integration = false,
     style = "default", -- or another Pygments style, e.g. "gruvbox-dark"
   })
 <

lua/pyrepl/config.lua

@@ -5,7 +5,7 @@ local defaults = {
     split_horizontal = false,
     split_ratio = 0.5,
     style = "default",
-    style_treesitter = true,
+    style_integration = true,
     image_max_history = 10,
     image_width_ratio = 0.5,
     image_height_ratio = 0.5,

lua/pyrepl/core.lua

@@ -96,7 +96,7 @@ local function open_new_repl(kernel)
     local console_path = python.get_console_path()
     local nvim_socket = vim.v.servername
     local style = config.get_state().style
-    local style_treesitter = config.get_state().style_treesitter
+    local style_integration = config.get_state().style_integration
 
     local buf = vim.api.nvim_create_buf(false, true)
     local win = open_scratch_win(buf)
@@ -113,11 +113,18 @@ local function open_new_repl(kernel)
         vim.o.termguicolors and "True" or "False",
     }
 
-    if style_treesitter then
-        local overrides = theme.build_pygments_theme()
-        if overrides then
+    if style_integration then
+        local pygments_overrides = theme.build_pygments_theme()
+        local prompt_toolkit_overrides = theme.build_prompt_toolkit_theme()
+
+        if pygments_overrides then
             cmd[#cmd + 1] = "--ZMQTerminalInteractiveShell.highlighting_style_overrides"
-            cmd[#cmd + 1] = overrides
+            cmd[#cmd + 1] = pygments_overrides
+        end
+
+        if prompt_toolkit_overrides then
+            cmd[#cmd + 1] = "--prompt-toolkit-overrides"
+            cmd[#cmd + 1] = prompt_toolkit_overrides
         end
     end
 

lua/pyrepl/theme.lua

@@ -56,31 +56,54 @@ local pygments_hl_map = {
     ["('OutPromptNum',)"] = { "@number", "Number" },
 }
 
----@param hl_name string
+local prompt_toolkit_hl_map = {
+    ["readline-like-completions"] = "Pmenu",
+    ["matching-bracket.other"] = "MatchParen",
+    ["matching-bracket.cursor"] = "Cursor",
+
+    ["completion-menu"] = "Pmenu",
+    ["completion-menu.completion"] = "Pmenu",
+    ["completion-menu.completion.current"] = "PmenuSel",
+    ["completion-menu.meta.completion"] = "PmenuExtra",
+    ["completion-menu.meta.completion.current"] = "PmenuExtraSel",
+    ["completion-menu.multi-column-meta"] = "PmenuExtra",
+
+    ["completion-toolbar"] = "Pmenu",
+    ["completion-toolbar.completion.current"] = "PmenuSel",
+}
+
+---@param hl string
 ---@return string|nil
-local function style_from_hl(hl_name)
-    local hl = vim.api.nvim_get_hl(0, {
-        name = hl_name,
+local function pygments_style_from_hl(hl)
+    local parts = {}
+    local style = vim.api.nvim_get_hl(0, {
+        name = hl,
         link = false,
     })
 
-    if type(hl.fg) ~= "number" then
-        return
+    if style.fg then
+        parts[#parts + 1] = string.format("fg:#%06x", style.fg)
+    end
+
+    if style.bg then
+        parts[#parts + 1] = string.format("bg:#%06x", style.bg)
     end
 
-    return string.format("'#%06x'", hl.fg)
+    if #parts > 0 then
+        return table.concat(parts, " ")
+    end
 end
 
 ---@return string|nil
 function M.build_pygments_theme()
     local theme = {}
 
-    for pygments, hls in pairs(pygments_hl_map) do
+    for pygments_hls, hls in pairs(pygments_hl_map) do
         -- obtain style from candidates
         for _, hl in ipairs(hls) do
-            local color = style_from_hl(hl)
-            if color then
-                theme[#theme + 1] = string.format("%s: %s", pygments, color)
+            local style = pygments_style_from_hl(hl)
+            if style then
+                theme[#theme + 1] = string.format("%s: '%s'", pygments_hls, style)
                 break
             end
         end
@@ -94,4 +117,21 @@ function M.build_pygments_theme()
     return "{" .. table.concat(theme, ", ") .. "}"
 end
 
+function M.build_prompt_toolkit_theme()
+    local theme = {}
+
+    for prompt_toolkit_hl, hl in pairs(prompt_toolkit_hl_map) do
+        local style = pygments_style_from_hl(hl)
+        if style then
+            theme[prompt_toolkit_hl] = style
+        end
+    end
+
+    if vim.tbl_isempty(theme) then
+        return nil
+    end
+
+    return vim.json.encode(theme)
+end
+
 return M

lua/pyrepl/types.lua

@@ -5,7 +5,7 @@
 ---@field split_horizontal boolean
 ---@field split_ratio number
 ---@field style string
----@field style_treesitter boolean
+---@field style_integration boolean
 ---@field image_max_history integer
 ---@field image_width_ratio number
 ---@field image_height_ratio number
@@ -20,7 +20,7 @@
 ---@field split_horizontal? boolean
 ---@field split_ratio? number
 ---@field style? string
----@field style_treesitter? boolean
+---@field style_integration? boolean
 ---@field image_max_history? integer
 ---@field image_width_ratio? number
 ---@field image_height_ratio? number

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "pyrepl"
-version = "0.6.5"
+version = "0.7.1"
 description = "REPL experience for Python in Neovim."
 readme = "README.md"
 authors = [

src/pyrepl/console.py

@@ -1,14 +1,16 @@
 import base64
 import io
+import json
 import os
-import sys
+from argparse import ArgumentParser
 from enum import StrEnum
 from queue import Queue
 from threading import Event, Thread
 from typing import Any
 
 import pynvim
 from jupyter_console.app import ZMQTerminalIPythonApp
+from prompt_toolkit.styles import defaults as ptk_defaults
 from traitlets.config import Config
 
 
@@ -64,7 +66,6 @@ def convert_image_to_png_base64(
             raw = image_data.encode("utf-8")
             png_bytes = cairosvg.svg2png(bytestring=raw)
             return base64.b64encode(png_bytes).decode("utf-8")
-
         except Exception:
             return None
 
@@ -153,9 +154,17 @@ def image_handler(data):
     config = Config()
     config.ZMQTerminalInteractiveShell.image_handler = "callable"
     config.ZMQTerminalInteractiveShell.callable_image_handler = image_handler
-
     app = ZMQTerminalIPythonApp.instance(config=config)
-    app.initialize(sys.argv[1:])
+
+    parser = ArgumentParser("Pyrepl console.")
+    parser.add_argument("--prompt-toolkit-overrides", type=str, default=None)
+    known, args = parser.parse_known_args()
+    app.initialize(args)
+
+    if known.prompt_toolkit_overrides is not None:
+        overrides = dict(ptk_defaults.PROMPT_TOOLKIT_STYLE)
+        overrides.update(json.loads(known.prompt_toolkit_overrides))
+        ptk_defaults.PROMPT_TOOLKIT_STYLE[:] = list(overrides.items())
 
     thread.start()
     app.start()  # type: ignore