diff --git a/AUTHORS b/AUTHORS index d0b4a5cee..481773ca8 100644 --- a/AUTHORS +++ b/AUTHORS @@ -16,6 +16,7 @@ Core Devs: Contributors: ------------- + * Muneer Alam (Muneer320) * AC (ac1982) * Wu Shuwen (dajiaohuang) * Brett diff --git a/changelog.rst b/changelog.rst index fc362092b..e22925c60 100644 --- a/changelog.rst +++ b/changelog.rst @@ -3,6 +3,8 @@ Upcoming Features: --------- +* Highlight function definitions from ``\sf`` and ``\sf+`` using the configured + syntax theme in psql, plain, and expanded output. * Add ``\\cd`` and ``\\ls`` filesystem commands, with path completion for filesystem-aware commands including ``\\e``, ``\\i``, ``\\log-file``, and ``\\o``. * Add a ``\service`` prompt token for connections made using ``service=...`` diff --git a/pgcli/main.py b/pgcli/main.py index 972b3c670..ce7cc1ec2 100644 --- a/pgcli/main.py +++ b/pgcli/main.py @@ -51,6 +51,8 @@ from prompt_toolkit.auto_suggest import AutoSuggestFromHistory from prompt_toolkit.cursor_shapes import ModalCursorShapeConfig from pygments.lexers.sql import PostgresLexer +from pygments import highlight +from pygments.formatters import Terminal256Formatter from pgspecial.main import PGSpecial, NO_QUERY, PAGER_OFF, PAGER_LONG_OUTPUT import pgspecial as special @@ -1421,7 +1423,15 @@ def _evaluate_command(self, text): title = None execution = time() - start - formatted = format_output(title, cur, headers, status, settings, self.explain_mode) + formatted = format_output( + title, + cur, + headers, + status, + settings, + self.explain_mode, + function_definition=bool(is_special and re.match(r"^\\sf\+?(?:\s|$)", sql.lstrip())), + ) output.extend(formatted) total = time() - start @@ -2158,7 +2168,7 @@ def exception_formatter(e, verbose_errors: bool = False): return click.style(s, fg="red") -def format_output(title, cur, headers, status, settings, explain_mode=False): +def format_output(title, cur, headers, status, settings, explain_mode=False, function_definition=False): output = [] expanded = settings.expanded or settings.table_format == "vertical" if settings.tuples_only: @@ -2191,6 +2201,20 @@ def format_arrays(data, headers, **_): return data, headers + def highlight_function(data, headers, table_format, **_): + lexer = PostgresLexer(stripnl=False, ensurenl=False) + sql_formatter = Terminal256Formatter(style=settings.style_output) + + def highlight_source(source): + colored = highlight(source, lexer, sql_formatter) + if table_format == "vertical": + # Tables normalize line endings; expanded output preserves the source. + line_endings = iter(re.findall(r"\r\n|\r|\n", source)) + colored = re.sub("\n", lambda _: next(line_endings), colored) + return colored + + return ([highlight_source(field) for field in row] for row in data), headers + def format_status(cur, status): # redshift does not return rowcount as part of status. # See https://github.com/dbcli/pgcli/issues/1320 @@ -2219,6 +2243,9 @@ def format_status(cur, status): if settings.column_date_formats: output_kwargs["preprocessors"] += (format_timestamps,) + if function_definition and settings.style_output and table_format in ("psql", "plain", "vertical"): + output_kwargs["preprocessors"] += (highlight_function,) + if table_format == "csv": # The default CSV dialect is "excel" which is not handling newline values correctly # Nevertheless, we want to keep on using "excel" on Windows since it uses '\r\n' diff --git a/tests/test_main.py b/tests/test_main.py index 57ba8afe8..5a2ec4205 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -32,6 +32,8 @@ from pgspecial.main import PAGER_OFF, PAGER_LONG_OUTPUT, PAGER_ALWAYS from utils import dbtest, run from collections import namedtuple +from cli_helpers.utils import strip_ansi +from pgcli.pgstyle import style_factory_output @pytest.mark.skipif(platform.system() == "Windows", reason="Not applicable in windows") @@ -81,6 +83,46 @@ def test_format_output(): assert list(results) == expected +@pytest.mark.parametrize("table_format", ["psql", "plain", "vertical"]) +@pytest.mark.parametrize("newline", ["\n", "\r\n", "\r"]) +def test_function_definition_highlighting(table_format, newline): + source = f"CREATE FUNCTION demo() RETURNS integer{newline}LANGUAGE sql AS 'SELECT 1';{newline}" + settings = OutputSettings(table_format=table_format, style_output=style_factory_output("native", {})) + plain = list(format_output(None, [(source,)], ["Source"], None, settings)) + colored = list(format_output(None, [(source,)], ["Source"], None, settings, function_definition=True)) + assert colored != plain + assert [strip_ansi(line) for line in colored] == [strip_ansi(line) for line in plain] + + +@pytest.mark.parametrize("table_format", ["csv", "tsv", "jsonl"]) +def test_function_definition_machine_output(table_format): + source = "CREATE FUNCTION demo() RETURNS text AS 'SELECT ''a,b''';\n" + settings = OutputSettings(table_format=table_format, style_output=style_factory_output("native", {})) + plain = list(format_output(None, [(source,)], ["Source"], None, settings)) + assert list(format_output(None, [(source,)], ["Source"], None, settings, function_definition=True)) == plain + + +def test_function_definition_auto_expand_preserves_line_endings(): + source = "CREATE FUNCTION demo()\r\nRETURNS integer LANGUAGE sql AS 'SELECT 1';\r\n" + settings = OutputSettings(table_format="psql", max_width=1, style_output=style_factory_output("native", {})) + plain = list(format_output(None, [(source,)], ["Source"], None, settings)) + colored = list(format_output(None, [(source,)], ["Source"], None, settings, function_definition=True)) + assert [strip_ansi(line) for line in colored] == [strip_ansi(line) for line in plain] + + +@pytest.mark.parametrize("command,highlighted", [(r"\sf demo", True), (r"\sf+ demo", True), (r"\sfoo demo", False)]) +def test_function_definition_command(command, highlighted): + source = "CREATE FUNCTION demo() RETURNS integer LANGUAGE sql AS 'SELECT 1';\n" + executor = mock.MagicMock() + executor.run.return_value = [(None, [(source,)], ["Source"], None, command, True, True)] + cli = PGCli(pgexecute=executor) + output, _ = cli._evaluate_command(command) + settings = OutputSettings(table_format=cli.table_format, style_output=cli.style_output) + plain = list(format_output(None, [(source,)], ["Source"], None, settings)) + assert (output != plain) == highlighted + assert [strip_ansi(line) for line in output] == [strip_ansi(line) for line in plain] + + def test_column_date_formats(): settings = OutputSettings( table_format="psql",