From d721e919e48d81f1717bd1e62cf28cc360a84547 Mon Sep 17 00:00:00 2001
From: Saugat Pachhai <suagatchhetri@outlook.com>
Date: Fri, 20 Nov 2020 16:38:32 +0545
Subject: [PATCH] parametrization: improve key error message (#4922)

---
 dvc/parsing/context.py            | 23 +++++++++++++++++++++--
 tests/func/test_stage_resolver.py |  4 ++--
 tests/unit/test_context.py        | 15 +++++++++++----
 3 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/dvc/parsing/context.py b/dvc/parsing/context.py
index 87b71834b..3432eacf5 100644
--- a/dvc/parsing/context.py
+++ b/dvc/parsing/context.py
@@ -45,6 +45,12 @@ class MergeError(ContextError):
         )
 
 
+class KeyNotInContext(ContextError):
+    def __init__(self, key) -> None:
+        self.key = key
+        super().__init__(f"Could not find '{key}'")
+
+
 def _merge(into, update, overwrite):
     for key, val in update.items():
         if isinstance(into.get(key), Mapping) and isinstance(val, Mapping):
@@ -171,7 +177,16 @@ class Container(Node, ABC):
             raise ValueError(
                 f"Could not find '{index}' in {self.data}"
             ) from exc
-        return d.select(rems[0]) if rems else d
+
+        if not rems:
+            return d
+
+        rem = rems[0]
+        if not isinstance(d, Container):
+            raise ValueError(
+                f"{index} is a primitive value, cannot get '{rem}'"
+            )
+        return d.select(rem)
 
     def get_sources(self):
         return {}
@@ -265,7 +280,11 @@ class Context(CtxDict):
                     Defaults to False. Note that the default is different from
                     `resolve`.
         """
-        node = super().select(key)
+        try:
+            node = super().select(key)
+        except ValueError as exc:
+            raise KeyNotInContext(key) from exc
+
         self._track_data(node)
         assert isinstance(node, Node)
         return node.value if unwrap else node
diff --git a/tests/func/test_stage_resolver.py b/tests/func/test_stage_resolver.py
index 860ed1076..c6377d68a 100644
--- a/tests/func/test_stage_resolver.py
+++ b/tests/func/test_stage_resolver.py
@@ -6,7 +6,7 @@ import pytest
 
 from dvc.dependency import _merge_params
 from dvc.parsing import DEFAULT_PARAMS_FILE, DataResolver
-from dvc.parsing.context import Node
+from dvc.parsing.context import KeyNotInContext, Node
 from dvc.path_info import PathInfo
 from dvc.utils.serialize import dump_json, dump_yaml
 
@@ -91,7 +91,7 @@ def test_no_params_yaml_and_vars(tmp_dir, dvc):
     resolver = DataResolver(
         dvc, PathInfo(str(tmp_dir)), deepcopy(TEMPLATED_DVC_YAML_DATA)
     )
-    with pytest.raises(ValueError):
+    with pytest.raises(KeyNotInContext):
         resolver.resolve()
 
 
diff --git a/tests/unit/test_context.py b/tests/unit/test_context.py
index 52be26b18..777f59d10 100644
--- a/tests/unit/test_context.py
+++ b/tests/unit/test_context.py
@@ -4,7 +4,14 @@ from math import pi
 import pytest
 
 from dvc.parsing import DEFAULT_PARAMS_FILE
-from dvc.parsing.context import Context, CtxDict, CtxList, MergeError, Value
+from dvc.parsing.context import (
+    Context,
+    CtxDict,
+    CtxList,
+    KeyNotInContext,
+    MergeError,
+    Value,
+)
 from dvc.tree.local import LocalTree
 from dvc.utils.serialize import dump_yaml
 from tests.func.test_stage_resolver import recurse_not_a_node
@@ -125,7 +132,7 @@ def test_select():
     assert context.select("lst") == CtxList([1, 2, 3])
     assert context.select("lst.0") == Value(1)
 
-    with pytest.raises(ValueError):
+    with pytest.raises(KeyNotInContext):
         context.select("baz")
 
     d = {
@@ -139,7 +146,7 @@ def test_select():
     assert context.select("lst.0") == CtxDict(d["lst"][0])
     assert context.select("lst.1") == CtxDict(d["lst"][1])
 
-    with pytest.raises(ValueError):
+    with pytest.raises(KeyNotInContext):
         context.select("lst.2")
 
     for i, _ in enumerate(d["lst"]):
@@ -247,7 +254,7 @@ def test_clone():
     assert c1 != c2
     assert c1 == Context(d)
     assert c2.select("lst.0.foo0") == Value("foo")
-    with pytest.raises(ValueError):
+    with pytest.raises(KeyNotInContext):
         c2.select("lst.1.foo1")