From 60be0427da9851914b73ce1b85e814cb7d2d56d0 Mon Sep 17 00:00:00 2001
From: Stephen Breakey <stephen.breakey@rank.com>
Date: Sat, 5 Dec 2020 00:43:08 +0200
Subject: [PATCH 1/2] Use pygment HTML() to style prompt

---
 mssqlcli/mssql_cli.py  | 16 ++++++++--------
 mssqlcli/mssqlclirc    |  8 ++++++++
 mssqlcli/mssqlstyle.py |  6 ++++++
 3 files changed, 22 insertions(+), 8 deletions(-)

diff --git a/mssqlcli/mssql_cli.py b/mssqlcli/mssql_cli.py
index 1a54badb..a94bd859 100644
--- a/mssqlcli/mssql_cli.py
+++ b/mssqlcli/mssql_cli.py
@@ -16,6 +16,7 @@
                                                       format_numbers)
 import humanize
 import click
+from prompt_toolkit import HTML
 from prompt_toolkit.shortcuts import PromptSession, CompleteStyle
 from prompt_toolkit.completion import DynamicCompleter, ThreadedCompleter
 from prompt_toolkit.enums import DEFAULT_BUFFER, EditingMode
@@ -457,8 +458,7 @@ def _build_cli(self, history):
         """
 
         def get_message():
-            prompt = self.get_prompt(self.prompt_format)
-            return [(u'class:prompt', prompt)]
+            return self.get_prompt(self.prompt_format)
 
         def get_continuation(width, line_number, is_soft_wrap):
             """
@@ -709,13 +709,13 @@ def get_completions(self, text, cursor_position):
                 Document(text=text, cursor_position=cursor_position), None)
 
     def get_prompt(self, string):
-        string = string.replace('\\t', self.now.strftime('%x %X'))
-        string = string.replace('\\u', self.mssqlcliclient_main.user_name or '(none)')
-        string = string.replace('\\h', self.mssqlcliclient_main.prompt_host or '(none)')
-        string = string.replace('\\d', self.mssqlcliclient_main.connected_database or '(none)')
-        string = string.replace('\\p', str(self.mssqlcliclient_main.prompt_port) or '(none)')
+        string = string.replace('\\t', "<prompt.datetime>%s</prompt.datetime>" % self.now.strftime('%x %X'))
+        string = string.replace('\\u', "<prompt.username>%s</prompt.username>" % self.mssqlcliclient_main.user_name or '(none)')
+        string = string.replace('\\h', "<prompt.hostname>%s</prompt.hostname>" % self.mssqlcliclient_main.prompt_host or '(none)')
+        string = string.replace('\\d', "<prompt.database>%s</prompt.database>" % self.mssqlcliclient_main.connected_database or '(none)')
+        string = string.replace('\\p', "<prompt.port>%s</prompt.port>" % str(self.mssqlcliclient_main.prompt_port) or '(none)')
         string = string.replace('\\n', "\n")
-        return string
+        return HTML("<prompt.default>%s</prompt.default>" % string)
 
     def get_last_query(self):
         """Get the last query executed or None."""
diff --git a/mssqlcli/mssqlclirc b/mssqlcli/mssqlclirc
index fc8deeb5..99e585c1 100644
--- a/mssqlcli/mssqlclirc
+++ b/mssqlcli/mssqlclirc
@@ -145,6 +145,14 @@ arg-toolbar.text = 'nobold'
 bottom-toolbar.transaction.valid = 'bg:#222222 #00ff5f bold'
 bottom-toolbar.transaction.failed = 'bg:#222222 #ff005f bold'
 
+# style classes for the prompt
+# prompt.default = 'bg:#000000 #FFFFFF'
+# prompt.datetime = 'bg:#000000 #FFFFFF'
+# prompt.username = 'bg:#000000 #FFFFFF'
+# prompt.hostname = 'bg:#000000 #FFFFFF'
+# prompt.database = 'bg:#000000 #FFFFFF'
+# prompt.port = 'bg:#000000 #FFFFFF'
+
 # style classes for colored table output
 output.header = "#00ff5f bold"
 output.odd-row = ""
diff --git a/mssqlcli/mssqlstyle.py b/mssqlcli/mssqlstyle.py
index 3f3ddb19..7a1303c7 100644
--- a/mssqlcli/mssqlstyle.py
+++ b/mssqlcli/mssqlstyle.py
@@ -36,6 +36,12 @@
     Token.Output.Header: 'output.header',
     Token.Output.OddRow: 'output.odd-row',
     Token.Output.EvenRow: 'output.even-row',
+    Token.Prompt.Default: 'prompt.default',
+    Token.Prompt.Datetime: 'prompt.datetime',
+    Token.Prompt.Username: 'prompt.username',
+    Token.Prompt.Hostname: 'prompt.hostname',
+    Token.Prompt.Database: 'prompt.database',
+    Token.Prompt.Port: 'prompt.port',
 }
 
 # reverse dict for cli_helpers, because they still expect Pygments tokens.

From 80b1643b4858b4f31c87274ccbb3a5e114c2fc0d Mon Sep 17 00:00:00 2001
From: Stephen Breakey <stephen.breakey@rank.com>
Date: Sat, 5 Dec 2020 00:55:36 +0200
Subject: [PATCH 2/2] Removed unnecessary TOKEN_TO_PROMPT_STYLE changes

---
 mssqlcli/mssqlstyle.py | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/mssqlcli/mssqlstyle.py b/mssqlcli/mssqlstyle.py
index 7a1303c7..3f3ddb19 100644
--- a/mssqlcli/mssqlstyle.py
+++ b/mssqlcli/mssqlstyle.py
@@ -36,12 +36,6 @@
     Token.Output.Header: 'output.header',
     Token.Output.OddRow: 'output.odd-row',
     Token.Output.EvenRow: 'output.even-row',
-    Token.Prompt.Default: 'prompt.default',
-    Token.Prompt.Datetime: 'prompt.datetime',
-    Token.Prompt.Username: 'prompt.username',
-    Token.Prompt.Hostname: 'prompt.hostname',
-    Token.Prompt.Database: 'prompt.database',
-    Token.Prompt.Port: 'prompt.port',
 }
 
 # reverse dict for cli_helpers, because they still expect Pygments tokens.