parametrization: print pretty interpolation errors (#4953)

* parametrization: print pretty interpolation errors

* Fix col/loc in exception msg

* Fix tests

dvc/parsing/interpolate.py

@@ -10,6 +10,7 @@ from pyparsing import (
 )
 
 from dvc.exceptions import DvcException
+from dvc.utils import colorize
 
 if typing.TYPE_CHECKING:
     from typing import List, Match
@@ -53,16 +54,31 @@ def is_interpolated_string(val):
 
 
 def format_and_raise_parse_error(exc):
-    msg = ParseException.explain(exc, depth=0)
-    raise ParseError(msg)
+    raise ParseError(_format_exc_msg(exc))
+
+
+def _format_exc_msg(exc: ParseException):
+    exc.loc += 2  # 2 because we append `${` at the start of expr below
+
+    expr = exc.pstr
+    exc.pstr = "${" + exc.pstr + "}"
+    error = ParseException.explain(exc, depth=0)
+
+    _, pointer, *explains = error.splitlines()
+    pstr = "{brace_open}{expr}{brace_close}".format(
+        brace_open=colorize("${", color="blue"),
+        expr=colorize(expr, color="magenta"),
+        brace_close=colorize("}", color="blue"),
+    )
+    msg = "\n".join(explains)
+    pointer = colorize(pointer, color="red")
+    return "\n".join([pstr, pointer, colorize(msg, color="red", style="bold")])
 
 
 def parse_expr(s: str):
     try:
         result = parser.parseString(s, parseAll=True)
     except ParseException as exc:
-        exc.pstr = "${" + exc.pstr + "}"
-        exc.loc += 2
         format_and_raise_parse_error(exc)
 
     joined = result.asList()

dvc/utils/__init__.py

@@ -236,20 +236,29 @@ def current_timestamp():
     return int(nanotime.timestamp(time.time()))
 
 
-def colorize(message, color=None):
+def colorize(message, color=None, style=None):
     """Returns a message in a specified color."""
     if not color:
         return message
 
+    styles = {
+        "dim": colorama.Style.DIM,
+        "bold": colorama.Style.BRIGHT,
+    }
+
     colors = {
         "green": colorama.Fore.GREEN,
         "yellow": colorama.Fore.YELLOW,
         "blue": colorama.Fore.BLUE,
         "red": colorama.Fore.RED,
+        "magenta": colorama.Fore.MAGENTA,
     }
 
-    return "{color}{message}{nc}".format(
-        color=colors.get(color, ""), message=message, nc=colorama.Fore.RESET
+    return "{style}{color}{message}{reset}".format(
+        style=styles.get(style, ""),
+        color=colors.get(color, ""),
+        message=message,
+        reset=colorama.Style.RESET_ALL,
     )
 
 

tests/func/test_stage_resolver_error_msg.py

@@ -1,4 +1,5 @@
 import logging
+import re
 
 import pytest
 
@@ -9,6 +10,11 @@ from dvc.utils.serialize import dump_json, dump_yaml
 DATA = {"models": {"bar": "bar", "foo": "foo"}}
 
 
+def escape_ansi(line):
+    ansi_escape = re.compile(r"(\x9B|\x1B\[)[0-?]*[ -\/]*[@-~]")
+    return ansi_escape.sub("", line)
+
+
 @pytest.fixture
 def repo(tmp_dir, dvc):
     dump_yaml(tmp_dir / DEFAULT_PARAMS_FILE, DATA)
@@ -73,7 +79,7 @@ def test_failed_to_interpolate(tmp_dir, repo):
     with pytest.raises(ResolveError) as exc_info:
         resolver = DataResolver(repo, tmp_dir, d)
         resolver.resolve()
-    assert str(exc_info.value) == (
+    assert escape_ansi(str(exc_info.value)) == (
         "failed to parse 'stages.build.cmd' in 'dvc.yaml':\n"
         "${models.foo.}\n"
         "            ^\n"
@@ -170,7 +176,7 @@ def test_foreach_interpolation_key_error(tmp_dir, repo):
     resolver = DataResolver(repo, tmp_dir, d)
     with pytest.raises(ResolveError) as exc_info:
         resolver.resolve()
-    assert str(exc_info.value) == (
+    assert escape_ansi(str(exc_info.value)) == (
         "failed to parse 'stages.build.foreach' in 'dvc.yaml':\n"
         "${models[123}\n"
         "        ^\n"
@@ -201,7 +207,7 @@ def test_foreach_generated_errors(tmp_dir, repo, syn, msg):
     resolver = DataResolver(repo, tmp_dir, d)
     with pytest.raises(ResolveError) as exc_info:
         resolver.resolve()
-    assert str(exc_info.value) == (
+    assert escape_ansi(str(exc_info.value)) == (
         "failed to parse '[email protected]' in 'dvc.yaml':" + msg
     )
 
@@ -260,7 +266,7 @@ def test_foreach_wdir_interpolation_issues(tmp_dir, repo, wdir, msg):
     resolver = DataResolver(repo, tmp_dir, d)
     with pytest.raises(ResolveError) as exc_info:
         resolver.resolve()
-    assert str(exc_info.value) == (
+    assert escape_ansi(str(exc_info.value)) == (
         "failed to parse '[email protected]' in 'dvc.yaml':" + msg
     )
 
@@ -297,7 +303,7 @@ def test_wdir_failed_to_interpolate(tmp_dir, repo, wdir, expected_msg):
     resolver = DataResolver(repo, tmp_dir, d)
     with pytest.raises(ResolveError) as exc_info:
         resolver.resolve()
-    assert str(exc_info.value) == (
+    assert escape_ansi(str(exc_info.value)) == (
         "failed to parse 'stages.build.wdir' in 'dvc.yaml':" + expected_msg
     )