Highlight function definitions from sf commands

AUTHORS

@@ -16,6 +16,7 @@ Core Devs:
 
 Contributors:
 -------------
+    * Muneer Alam (Muneer320)
     * AC (ac1982)
     * Wu Shuwen (dajiaohuang)
     * Brett

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=...``

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'

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",