parse: parse target correctly for generated stages (iterative#4961)

9173860 · Nov 25, 2020 · 1294450 · 1294450
1 parent c989b28
commit 1294450
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 2 deletions.
diff --git a/dvc/utils/__init__.py b/dvc/utils/__init__.py
@@ -409,21 +409,31 @@ def error_link(name):
     return format_link(f"https://error.dvc.org/{name}")
 
 
-def parse_target(target, default=None):
+def parse_target(target: str, default: str = None):
     from dvc.dvcfile import PIPELINE_FILE, PIPELINE_LOCK, is_valid_filename
     from dvc.exceptions import DvcException
+    from dvc.parsing import JOIN
 
     if not target:
         return None, None
 
-    match = TARGET_REGEX.match(target)
+    # look for first "@", so as not to assume too much about stage name
+    # eg: it might contain ":" in a generated stages from dict which might
+    # affect further parsings with the regex.
+    group, _, key = target.partition(JOIN)
+    match = TARGET_REGEX.match(group)
+
     if not match:
         return target, None
 
     path, name = (
         match.group("path"),
         match.group("name"),
     )
+
+    if name and key:
+        name += f"{JOIN}{key}"
+
     if path:
         if os.path.basename(path) == PIPELINE_LOCK:
             raise DvcException(

diff --git a/tests/unit/utils/test_utils.py b/tests/unit/utils/test_utils.py
@@ -149,6 +149,31 @@ def test_resolve_output(inp, out, is_dir, expected, mocker):
         ["something.dvc:name", ("something.dvc", "name"), None],
         ["../something.dvc:name", ("../something.dvc", "name"), None],
         ["file", (None, "file"), None],
+        ["build@15", (None, "build@15"), None],
+        ["build@{'level': 35}", (None, "build@{'level': 35}"), None],
+        [":build@15", ("dvc.yaml", "build@15"), None],
+        [":build@{'level': 35}", ("dvc.yaml", "build@{'level': 35}"), None],
+        ["dvc.yaml:build@15", ("dvc.yaml", "build@15"), None],
+        [
+            "dvc.yaml:build@{'level': 35}",
+            ("dvc.yaml", "build@{'level': 35}"),
+            None,
+        ],
+        [
+            "build2@{'level': [1, 2, 3]}",
+            (None, "build2@{'level': [1, 2, 3]}"),
+            None,
+        ],
+        [
+            ":build2@{'level': [1, 2, 3]}",
+            ("dvc.yaml", "build2@{'level': [1, 2, 3]}"),
+            None,
+        ],
+        [
+            "dvc.yaml:build2@{'level': [1, 2, 3]}",
+            ("dvc.yaml", "build2@{'level': [1, 2, 3]}"),
+            None,
+        ],
     ],
 )
 def test_parse_target(inp, out, default):