From cb3e52829b601ae574ce0178f58e605e7a37f2b3 Mon Sep 17 00:00:00 2001
From: Saugat Pachhai <suagatchhetri@outlook.com>
Date: Fri, 20 Nov 2020 14:31:16 +0545
Subject: [PATCH] interpolation: support `[]` indices syntax (#4918)

Via pyparsing. The error message is not that great, as it's greedy.
But, it does point out where the mistakes are, correctly most of the
time. :)

Also raises appropriate error messages for ParseError
---
 dvc/parsing/interpolate.py | 57 ++++++++++++++++++++++++++++++++++++--
 setup.py                   |  1 +
 2 files changed, 55 insertions(+), 3 deletions(-)

diff --git a/dvc/parsing/interpolate.py b/dvc/parsing/interpolate.py
index 18cb3a735..0bbe3f3e3 100644
--- a/dvc/parsing/interpolate.py
+++ b/dvc/parsing/interpolate.py
@@ -1,22 +1,48 @@
 import re
 import typing
 
+from pyparsing import (
+    CharsNotIn,
+    ParseException,
+    ParserElement,
+    Suppress,
+    ZeroOrMore,
+)
+
+from dvc.exceptions import DvcException
+
 if typing.TYPE_CHECKING:
     from typing import List, Match
 
     from .context import Context
 
+LBRACK = "["
+RBRACK = "]"
+PERIOD = "."
 KEYCRE = re.compile(
     r"""
     (?<!\\)                   # escape \${} or ${{}}
     \$                        # starts with $
     (?:({{)|({))              # either starts with double braces or single
-    ([\w._ \\/-]*?)           # match every char, attr access through "."
+    (.*?)                     # match every char inside
     (?(1)}})(?(2)})           # end with same kinds of braces it opened with
 """,
     re.VERBOSE,
 )
 
+ParserElement.enablePackrat()
+
+
+word = CharsNotIn(f"{PERIOD}{LBRACK}{RBRACK}")
+idx = Suppress(LBRACK) + word + Suppress(RBRACK)
+attr = Suppress(PERIOD) + word
+parser = word + ZeroOrMore(attr ^ idx)
+parser.setParseAction(PERIOD.join)
+
+
+class ParseError(DvcException):
+    pass
+
 
 def get_matches(template: str):
     return list(KEYCRE.finditer(template))
@@ -26,17 +52,42 @@ def is_interpolated_string(val):
     return bool(get_matches(val)) if isinstance(val, str) else False
 
 
+def format_and_raise_parse_error(exc):
+    msg = ParseException.explain(exc, depth=0)
+    raise ParseError(msg)
+
+
+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()
+    assert len(joined) == 1
+    return joined[0]
+
+
 def get_expression(match: "Match"):
     _, _, inner = match.groups()
-    return inner
+    return parse_expr(inner)
 
 
 def str_interpolate(template: str, matches: "List[Match]", context: "Context"):
+    from .context import PRIMITIVES
+
     index, buf = 0, ""
     for match in matches:
         start, end = match.span(0)
         expr = get_expression(match)
-        buf += template[index:start] + str(context.select(expr))
+        value = context.select(expr, unwrap=True)
+        if value is not None and not isinstance(value, PRIMITIVES):
+            raise ParseError(
+                f"Cannot interpolate data of type '{type(value).__name__}'"
+            )
+        buf += template[index:start] + str(value)
         index = end
     buf += template[index:]
     # regex already backtracks and avoids any `${` starting with
diff --git a/setup.py b/setup.py
index 98c58135f..8a5bf94b7 100644
--- a/setup.py
+++ b/setup.py
@@ -83,6 +83,7 @@ install_requires = [
     "rich>=3.0.5",
     "dictdiffer>=0.8.1",
     "python-benedict>=0.21.1",
+    "pyparsing==2.4.7",
 ]