Refactor context set/alias to Context class (#4831)

dvc/parsing/__init__.py

@@ -12,13 +12,7 @@ from dvc.dependency.param import ParamsDependency
 from dvc.path_info import PathInfo
 
 from .context import Context
-from .interpolate import (
-    _get_matches,
-    _is_exact_string,
-    _is_interpolated_string,
-    _resolve_str,
-    resolve,
-)
+from .interpolate import resolve
 
 if TYPE_CHECKING:
     from dvc.repo import Repo
@@ -157,43 +151,4 @@ class DataResolver:
     @classmethod
     def set_context_from(cls, context: Context, to_set):
         for key, value in to_set.items():
-            if key in context:
-                raise ValueError(f"Cannot set '{key}', key already exists")
-            if isinstance(value, str):
-                cls._check_joined_with_interpolation(key, value)
-                value = _resolve_str(value, context, unwrap=False)
-            elif isinstance(value, (Sequence, Mapping)):
-                cls._check_nested_collection(key, value)
-                cls._check_interpolation_collection(key, value)
-            context[key] = value
-
-    @staticmethod
-    def _check_nested_collection(key: str, value: SeqOrMap):
-        values = value.values() if isinstance(value, Mapping) else value
-        has_nested = any(
-            not isinstance(item, str) and isinstance(item, (Mapping, Sequence))
-            for item in values
-        )
-        if has_nested:
-            raise ValueError(f"Cannot set '{key}', has nested dict/list")
-
-    @staticmethod
-    def _check_interpolation_collection(key: str, value: SeqOrMap):
-        values = value.values() if isinstance(value, Mapping) else value
-        interpolated = any(_is_interpolated_string(item) for item in values)
-        if interpolated:
-            raise ValueError(
-                f"Cannot set '{key}', "
-                "having interpolation inside "
-                f"'{type(value).__name__}' is not supported."
-            )
-
-    @staticmethod
-    def _check_joined_with_interpolation(key: str, value: str):
-        matches = _get_matches(value)
-        if matches and not _is_exact_string(value, matches):
-            raise ValueError(
-                f"Cannot set '{key}', "
-                "joining string with interpolated string"
-                "is not supported"
-            )
+            context.set(key, value)

dvc/parsing/context.py

@@ -1,15 +1,23 @@
 import os
 from collections import defaultdict
-from collections.abc import Mapping, MutableMapping, MutableSequence
+from collections.abc import Mapping, MutableMapping, MutableSequence, Sequence
 from contextlib import contextmanager
 from copy import deepcopy
 from dataclasses import dataclass, field, replace
-from typing import Any, List, Optional, Sequence, Union
+from typing import Any, List, Optional, Union
 
 from funcy import identity
 
+from dvc.parsing.interpolate import (
+    get_matches,
+    is_exact_string,
+    is_interpolated_string,
+    resolve_str,
+)
 from dvc.utils.serialize import LOADERS
 
+SeqOrMap = Union[Sequence, Mapping]
+
 
 def _merge(into, update, overwrite):
     for key, val in update.items():
@@ -210,3 +218,60 @@ class Context(CtxDict):
     def clone(cls, ctx: "Context") -> "Context":
         """Clones given context."""
         return cls(deepcopy(ctx.data))
+
+    def set(self, key, value):
+        """
+        Sets a value, either non-interpolated values to a key,
+        or an interpolated string after resolving it.
+
+        >>> c = Context({"foo": "foo", "bar": [1, 2], "lorem": {"a": "z"}})
+        >>> c.set("foobar", "${bar}")
+        >>> c
+        {'foo': 'foo', 'bar': [1, 2], 'lorem': {'a': 'z'}, 'foobar': [1, 2]}
+        """
+        if key in self:
+            raise ValueError(f"Cannot set '{key}', key already exists")
+        if isinstance(value, str):
+            self._check_joined_with_interpolation(key, value)
+            value = resolve_str(value, self, unwrap=False)
+        elif isinstance(value, (Sequence, Mapping)):
+            self._check_not_nested_collection(key, value)
+            self._check_interpolation_collection(key, value)
+        self[key] = value
+
+    @staticmethod
+    def _check_not_nested_collection(key: str, value: SeqOrMap):
+        values = value.values() if isinstance(value, Mapping) else value
+        has_nested = any(
+            not isinstance(item, str) and isinstance(item, (Mapping, Sequence))
+            for item in values
+        )
+        if has_nested:
+            raise ValueError(f"Cannot set '{key}', has nested dict/list")
+
+    @staticmethod
+    def _check_interpolation_collection(key: str, value: SeqOrMap):
+        values = value.values() if isinstance(value, Mapping) else value
+        interpolated = any(is_interpolated_string(item) for item in values)
+        if interpolated:
+            raise ValueError(
+                f"Cannot set '{key}', "
+                "having interpolation inside "
+                f"'{type(value).__name__}' is not supported."
+            )
+
+    @staticmethod
+    def _check_joined_with_interpolation(key: str, value: str):
+        matches = get_matches(value)
+        if matches and not is_exact_string(value, matches):
+            raise ValueError(
+                f"Cannot set '{key}', "
+                "joining string with interpolated string"
+                "is not supported"
+            )
+
+
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()

dvc/parsing/interpolate.py

@@ -1,9 +1,11 @@
 import re
+import typing
 from collections.abc import Mapping
 
 from funcy import rpartial
 
-from dvc.parsing.context import Context, Value
+if typing.TYPE_CHECKING:
+    from .context import Context
 
 KEYCRE = re.compile(
     r"""
@@ -19,21 +21,23 @@ KEYCRE = re.compile(
 UNWRAP_DEFAULT = True
 
 
-def _get_matches(template: str):
+def get_matches(template: str):
     return list(KEYCRE.finditer(template))
 
 
-def _is_interpolated_string(val):
-    return bool(_get_matches(val)) if isinstance(val, str) else False
+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 _resolve_value(match, context: "Context", unwrap=UNWRAP_DEFAULT):
     _, _, inner = match.groups()
     value = context.select(inner)
     return _unwrap(value) if unwrap else value
@@ -48,13 +52,13 @@ def _str_interpolate(template, matches, context):
     return buf + template[index:]
 
 
-def _is_exact_string(src: str, matches):
+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):
+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)
@@ -75,5 +79,5 @@ def resolve(src, context):
     elif isinstance(src, Seq):
         return type(src)(map(apply_value, src))
     elif isinstance(src, str):
-        return _resolve_str(src, context)
+        return resolve_str(src, context)
     return src