Allow the argument to /prompt to be quoted

changelog.md

@@ -12,6 +12,7 @@ Features
 * Ability to set the `prompt` value as a DSN query parameter.
 * Add `--more` option to some `/dsn` subcommands, showing more DSN query parameters.
 * Let `/prompt` with no arguments display the current prompt format string.
+* Allow the `/prompt` argument to be quoted.
 
 
 Bugfixes

mycli/client_commands.py

@@ -166,6 +166,9 @@ def change_prompt_format(self, arg: str, **_) -> list[SQLResult]:
         if not arg:
             return [SQLResult(status=f'Prompt format: "{self.prompt_format}"')]
 
+        if len(arg) >= 2 and arg[0] == arg[-1] and arg[0] in {'\'', '"'}:
+            arg = arg[1:-1]
+
         self.prompt_format = arg
         return [SQLResult(status=f'Changed prompt format to: "{arg}"')]
 

test/pytests/test_client_commands.py

@@ -8,6 +8,8 @@
 
 from mycli import client_commands
 from mycli.client_commands import ClientCommandsMixin
+from mycli.packages import special
+from mycli.packages.special import main as special_main
 from mycli.packages.sqlresult import SQLResult
 
 
@@ -240,6 +242,51 @@ def test_change_prompt_format_updates_prompt_format() -> None:
     assert client.prompt_format == '\\u> '
 
 
+@pytest.mark.parametrize(
+    ('command', 'quote'),
+    [
+        ('/prompt', '"'),
+        (r'\R', "'"),
+    ],
+)
+def test_change_prompt_format_accepts_quoted_value(
+    monkeypatch: pytest.MonkeyPatch,
+    command: str,
+    quote: str,
+) -> None:
+    monkeypatch.setattr(special_main, 'COMMANDS', {})
+    monkeypatch.setattr(special_main, 'CASE_SENSITIVE_COMMANDS', set())
+    monkeypatch.setattr(special_main, 'CASE_INSENSITIVE_COMMANDS', set())
+    client = DummyClient()
+    client.prompt_format = 'old> '
+    client.register_special_commands()
+
+    result = special.execute(None, f'{command} {quote} \\u> {quote}')
+
+    assert result == [SQLResult(status='Changed prompt format to: " \\u> "')]
+    assert client.prompt_format == ' \\u> '
+
+
+@pytest.mark.parametrize('arg', ["''", '""'])
+def test_change_prompt_format_accepts_quoted_empty_value(arg: str) -> None:
+    client = DummyClient()
+    client.prompt_format = 'old> '
+
+    result = client.change_prompt_format(arg)
+
+    assert result == [SQLResult(status='Changed prompt format to: ""')]
+    assert client.prompt_format == ''
+
+
+@pytest.mark.parametrize('arg', ["'unmatched", '"unmatched'])
+def test_change_prompt_format_keeps_unmatched_quote(arg: str) -> None:
+    client = DummyClient()
+
+    client.change_prompt_format(arg)
+
+    assert client.prompt_format == arg
+
+
 def test_initialize_logging_uses_null_handler_for_none_level(monkeypatch: pytest.MonkeyPatch) -> None:
     client = DummyClient()
     client.config = {'main': {'log_file': '/unused/mycli.log', 'log_level': 'NONE'}}