From e9a9415fbf84e15eb26b8c2f19caffd79417bdc3 Mon Sep 17 00:00:00 2001
From: Saugat Pachhai <suagatchhetri@outlook.com>
Date: Wed, 4 Nov 2020 17:04:40 +0545
Subject: [PATCH] refactor: move interpolation logic inside the Context (#4840)

* Move interpolation logic to the Context

* Don't unwrap by default on `Context.select`
---
 dvc/parsing/__init__.py    |  5 ++---
 dvc/parsing/context.py     | 22 ++++++++++++++-----
 dvc/parsing/interpolate.py | 45 +++++++++++---------------------------
 3 files changed, 31 insertions(+), 41 deletions(-)

diff --git a/dvc/parsing/__init__.py b/dvc/parsing/__init__.py
index 5fffb93be..1a8e42fe4 100644
--- a/dvc/parsing/__init__.py
+++ b/dvc/parsing/__init__.py
@@ -4,7 +4,7 @@ from collections import defaultdict
 from collections.abc import Mapping, Sequence
 from copy import deepcopy
 from itertools import starmap
-from typing import TYPE_CHECKING, Union
+from typing import TYPE_CHECKING
 
 from funcy import first, join
 
@@ -29,7 +29,6 @@ IN_KWD = "in"
 SET_KWD = "set"
 
 DEFAULT_SENTINEL = object()
-SeqOrMap = Union[Sequence, Mapping]
 
 
 class DataResolver:
@@ -125,7 +124,7 @@ class DataResolver:
         if not wdir:
             return self.wdir
 
-        wdir = context.resolve(wdir)
+        wdir = str(context.resolve_str(wdir, unwrap=True))
         return self.wdir / str(wdir)
 
     def _foreach(self, context: Context, name: str, foreach_data, in_data):
diff --git a/dvc/parsing/context.py b/dvc/parsing/context.py
index 1c94f7155..ae35ca80d 100644
--- a/dvc/parsing/context.py
+++ b/dvc/parsing/context.py
@@ -9,11 +9,11 @@ from typing import Any, List, Optional, Union
 from funcy import identity
 
 from dvc.parsing.interpolate import (
-    UNWRAP_DEFAULT,
+    get_expression,
     get_matches,
     is_exact_string,
     is_interpolated_string,
-    resolve_str,
+    str_interpolate,
 )
 from dvc.utils.serialize import LOADERS
 
@@ -202,10 +202,13 @@ class Context(CtxDict):
     def tracked(self):
         return self._tracked_data
 
-    def select(self, key: str):
+    def select(
+        self, key: str, unwrap=False
+    ):  # pylint: disable=arguments-differ
+        # NOTE: `unwrap` default value is different from `resolve_str`
         node = super().select(key)
         self._track_data(node)
-        return node
+        return node.value if isinstance(node, Value) and unwrap else node
 
     @classmethod
     def load_from(cls, tree, file: str) -> "Context":
@@ -256,7 +259,7 @@ class Context(CtxDict):
             return self.resolve_str(src)
         return src
 
-    def resolve_str(self, src, unwrap=UNWRAP_DEFAULT):
+    def resolve_str(self, src: str, unwrap=True):
         """Resolves interpolated string to it's original value,
         or in case of multiple interpolations, a combined string.
 
@@ -266,7 +269,14 @@ class Context(CtxDict):
         >>> c.resolve_str("enabled? ${enabled}")
         'enabled? True'
         """
-        return resolve_str(src, self, unwrap=unwrap)
+        matches = get_matches(src)
+        if is_exact_string(src, matches):
+            # replace "${enabled}", if `enabled` is a boolean, with it's actual
+            # value rather than it's string counterparts.
+            expr = get_expression(matches[0])
+            return self.select(expr, unwrap=unwrap)
+        # but not "${num} days"
+        return str_interpolate(src, matches, self)
 
     @staticmethod
     def _check_not_nested_collection(key: str, value: SeqOrMap):
diff --git a/dvc/parsing/interpolate.py b/dvc/parsing/interpolate.py
index f122bf66a..18cb3a735 100644
--- a/dvc/parsing/interpolate.py
+++ b/dvc/parsing/interpolate.py
@@ -2,6 +2,8 @@ import re
 import typing
 
 if typing.TYPE_CHECKING:
+    from typing import List, Match
+
     from .context import Context
 
 KEYCRE = re.compile(
@@ -15,8 +17,6 @@ KEYCRE = re.compile(
     re.VERBOSE,
 )
 
-UNWRAP_DEFAULT = True
-
 
 def get_matches(template: str):
     return list(KEYCRE.finditer(template))
@@ -26,42 +26,23 @@ def is_interpolated_string(val):
     return bool(get_matches(val)) if isinstance(val, str) else False
 
 
-def _unwrap(value):
-    from .context import Value
-
-    if isinstance(value, Value):
-        return value.value
-    return value
-
-
-def _resolve_value(match, context: "Context", unwrap=UNWRAP_DEFAULT):
+def get_expression(match: "Match"):
     _, _, inner = match.groups()
-    value = context.select(inner)
-    return _unwrap(value) if unwrap else value
+    return inner
 
 
-def _str_interpolate(template, matches, context):
+def str_interpolate(template: str, matches: "List[Match]", context: "Context"):
     index, buf = 0, ""
     for match in matches:
         start, end = match.span(0)
-        buf += template[index:start] + str(_resolve_value(match, context))
+        expr = get_expression(match)
+        buf += template[index:start] + str(context.select(expr))
         index = end
-    return buf + template[index:]
-
-
-def is_exact_string(src: str, matches):
-    return len(matches) == 1 and src == matches[0].group(0)
-
-
-def resolve_str(src: str, context, unwrap=UNWRAP_DEFAULT):
-    matches = get_matches(src)
-    if is_exact_string(src, matches):
-        # replace "${enabled}", if `enabled` is a boolean, with it's actual
-        # value rather than it's string counterparts.
-        return _resolve_value(matches[0], context, unwrap=unwrap)
-
-    # but not "${num} days"
-    src = _str_interpolate(src, matches, context)
+    buf += template[index:]
     # regex already backtracks and avoids any `${` starting with
     # backslashes(`\`). We just need to replace those by `${`.
-    return src.replace(r"\${", "${")
+    return buf.replace(r"\${", "${")
+
+
+def is_exact_string(src: str, matches: "List[Match]"):
+    return len(matches) == 1 and src == matches[0].group(0)