From a38941b47e6db7d7b9cf0a77b92bf6c6ef22f029 Mon Sep 17 00:00:00 2001 From: Saugat Pachhai Date: Tue, 3 Nov 2020 17:06:46 +0545 Subject: [PATCH] Refactor context set/alias to Context class (#4831) --- dvc/parsing/__init__.py | 49 ++------------------------- dvc/parsing/context.py | 69 ++++++++++++++++++++++++++++++++++++-- dvc/parsing/interpolate.py | 24 +++++++------ 3 files changed, 83 insertions(+), 59 deletions(-) diff --git a/dvc/parsing/__init__.py b/dvc/parsing/__init__.py index 611390659..1ce85c24f 100644 --- a/dvc/parsing/__init__.py +++ b/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) diff --git a/dvc/parsing/context.py b/dvc/parsing/context.py index fd192fcee..994e2138a 100644 --- a/dvc/parsing/context.py +++ b/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() diff --git a/dvc/parsing/interpolate.py b/dvc/parsing/interpolate.py index 51160fec7..3be15ed49 100644 --- a/dvc/parsing/interpolate.py +++ b/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