move parsing params to utils (#5259)

* move parsing params to utils

* fix tests

* move chunk_dict and test it

* update test

dvc/repo/experiments/run.py

@@ -1,43 +1,17 @@
 import logging
 from typing import Iterable, Optional
 
-from dvc.exceptions import InvalidArgumentError
 from dvc.repo import locked
+from dvc.utils.cli_parse import loads_params_from_cli
 
 logger = logging.getLogger(__name__)
 
 
-def _parse_params(path_params: Iterable):
-    from ruamel.yaml import YAMLError
-
-    from dvc.dependency.param import ParamsDependency
-    from dvc.utils.serialize import loads_yaml
-
-    ret = {}
-    for path_param in path_params:
-        path, _, params_str = path_param.rpartition(":")
-        # remove empty strings from params, on condition such as `-p "file1:"`
-        params = {}
-        for param_str in filter(bool, params_str.split(",")):
-            try:
-                # interpret value strings using YAML rules
-                key, value = param_str.split("=")
-                params[key] = loads_yaml(value)
-            except (ValueError, YAMLError):
-                raise InvalidArgumentError(
-                    f"Invalid param/value pair '{param_str}'"
-                )
-        if not path:
-            path = ParamsDependency.DEFAULT_PARAMS_FILE
-        ret[path] = params
-    return ret
-
-
 @locked
 def run(
     repo,
-    targets: Optional[Iterable] = None,
-    params: Optional[Iterable] = None,
+    targets: Optional[Iterable[str]] = None,
+    params: Optional[Iterable[str]] = None,
     run_all: Optional[bool] = False,
     jobs: Optional[int] = 1,
     tmp_dir: Optional[bool] = False,
@@ -54,9 +28,7 @@ def run(
         return repo.experiments.reproduce_queued(jobs=jobs)
 
     if params:
-        params = _parse_params(params)
-    else:
-        params = []
+        params = loads_params_from_cli(params)
     return repo.experiments.reproduce_one(
         targets=targets, params=params, tmp_dir=tmp_dir, **kwargs
     )

dvc/repo/run.py

@@ -1,7 +1,10 @@
 import os
 from contextlib import suppress
 
-from funcy import concat, first, lfilter, without
+from funcy import concat, first, without
+
+from dvc.utils.cli_parse import parse_params_from_cli
+from dvc.utils.collections import chunk_dict
 
 from . import locked
 from .scm_context import scm_context
@@ -13,19 +16,6 @@ def is_valid_name(name: str):
     return not INVALID_STAGENAME_CHARS & set(name)
 
 
-def parse_params(path_params):
-    ret = []
-    for path_param in path_params:
-        path, _, params_str = path_param.rpartition(":")
-        # remove empty strings from params, on condition such as `-p "file1:"`
-        params = lfilter(bool, params_str.split(","))
-        if not path:
-            ret.extend(params)
-        else:
-            ret.append({path: params})
-    return ret
-
-
 def _get_file_path(kwargs):
     from dvc.dvcfile import DVC_FILE, DVC_FILE_SUFFIX
 
@@ -109,7 +99,7 @@ def run(self, fname=None, no_exec=False, single_stage=False, **kwargs):
         if not is_valid_name(stage_name):
             raise InvalidStageName
 
-    params = parse_params(kwargs.pop("params", []))
+    params = chunk_dict(parse_params_from_cli(kwargs.pop("params", [])))
     stage = create_stage(
         stage_cls, repo=self, path=path, params=params, **kwargs
     )

dvc/utils/cli_parse.py

@@ -0,0 +1,45 @@
+from collections import defaultdict
+from typing import Any, Dict, Iterable, List
+
+
+def parse_params_from_cli(path_params: Iterable[str]) -> Dict[str, List[str]]:
+    """Normalizes the shape of params from the CLI to dict."""
+    from dvc.dependency.param import ParamsDependency
+
+    ret: Dict[str, List[str]] = defaultdict(list)
+    for path_param in path_params:
+        path, _, params_str = path_param.rpartition(":")
+        # remove empty strings from params, on condition such as `-p "file1:"`
+        params = filter(bool, params_str.split(","))
+        if not path:
+            path = ParamsDependency.DEFAULT_PARAMS_FILE
+        ret[path].extend(params)
+    return ret
+
+
+def loads_params_from_cli(
+    path_params: Iterable[str],
+) -> Dict[str, Dict[str, Any]]:
+
+    """Loads the content of params from the cli as Python object."""
+    from ruamel.yaml import YAMLError
+
+    from dvc.exceptions import InvalidArgumentError
+
+    from .serialize import loads_yaml
+
+    normalized_params = parse_params_from_cli(path_params)
+    ret: Dict[str, Dict[str, Any]] = defaultdict(dict)
+
+    for path, param_keys in normalized_params.items():
+        for param_str in param_keys:
+            try:
+                key, _, value = param_str.partition("=")
+                # interpret value strings using YAML rules
+                parsed = loads_yaml(value)
+                ret[path][key] = parsed
+            except (ValueError, YAMLError):
+                raise InvalidArgumentError(
+                    f"Invalid param/value pair '{param_str}'"
+                )
+    return ret

dvc/utils/collections.py

@@ -1,6 +1,6 @@
 import os
 from collections.abc import Mapping
-from typing import Iterable, List, Union
+from typing import Dict, Iterable, List, TypeVar, Union
 
 from pygtrie import StringTrie as _StringTrie
 
@@ -66,3 +66,13 @@ def ensure_list(item: Union[Iterable[str], str, None]) -> List[str]:
     if isinstance(item, str):
         return [item]
     return list(item)
+
+
+_KT = TypeVar("_KT")
+_VT = TypeVar("_VT")
+
+
+def chunk_dict(d: Dict[_KT, _VT], size: int = 1) -> List[Dict[_KT, _VT]]:
+    from funcy import chunks
+
+    return [{key: d[key] for key in chunk} for chunk in chunks(size, d)]

tests/unit/test_run.py

@@ -1,29 +1,6 @@
 import pytest
 
-from dvc.repo.run import is_valid_name, parse_params
-
-
-def test_parse_params():
-    assert parse_params(
-        [
-            "param1",
-            "file1:param1,param2",
-            "file2:param2",
-            "file1:param2,param3,",
-            "param1,param2",
-            "param3,",
-            "file3:",
-        ]
-    ) == [
-        "param1",
-        {"file1": ["param1", "param2"]},
-        {"file2": ["param2"]},
-        {"file1": ["param2", "param3"]},
-        "param1",
-        "param2",
-        "param3",
-        {"file3": []},
-    ]
+from dvc.repo.run import is_valid_name
 
 
 @pytest.mark.parametrize("name", ["copy_name", "copy-name", "copyName", "12"])

tests/unit/utils/test_cli_parse.py

@@ -0,0 +1,20 @@
+from dvc.utils.cli_parse import parse_params_from_cli
+
+
+def test_parse_params():
+    assert parse_params_from_cli(
+        [
+            "param1",
+            "file1:param1,param2",
+            "file2:param2",
+            "file1:param3,param4,",
+            "param2,param10",
+            "param3,",
+            "file3:",
+        ]
+    ) == {
+        "params.yaml": ["param1", "param2", "param10", "param3"],
+        "file1": ["param1", "param2", "param3", "param4"],
+        "file2": ["param2"],
+        "file3": [],
+    }

tests/unit/utils/test_collections.py

@@ -1,5 +1,5 @@
 # pylint: disable=unidiomatic-typecheck
-from dvc.utils.collections import apply_diff
+from dvc.utils.collections import apply_diff, chunk_dict
 
 
 class MyDict(dict):
@@ -48,3 +48,13 @@ def test_apply_diff_seq():
     dest = {"l": inner}
     apply_diff(src, dest)
     assert dest["l"] is inner, "Updates inner lists"
+
+
+def test_chunk_dict():
+    assert chunk_dict({}) == []
+
+    d = {"a": 1, "b": 2, "c": 3}
+    assert chunk_dict(d) == [{"a": 1}, {"b": 2}, {"c": 3}]
+    assert chunk_dict(d, 2) == [{"a": 1, "b": 2}, {"c": 3}]
+    assert chunk_dict(d, 3) == [d]
+    assert chunk_dict(d, 4) == [d]