From f51bef040a37deb3f1bca1ece5b1d208e163d31b Mon Sep 17 00:00:00 2001
From: Ioana Grigoropol <14314993+ju0gri@users.noreply.github.com>
Date: Sat, 19 Dec 2020 19:38:40 +0000
Subject: [PATCH] pull: add glob option (#5032)

* api: add glob option for pull command

Related: #4816

Signed-off-by: Ioana Grigoropol <ioana.grigoropol@gmail.com>

* api: add globbing utility function

Related: #4816

Signed-off-by: Ioana Grigoropol <ioana.grigoropol@gmail.com>

* api: use utility function for pull command

Signed-off-by: Ioana Grigoropol <ioana.grigoropol@gmail.com>

* Update dvc/utils/__init__.py

Co-authored-by: Ruslan Kuprieiev <kupruser@gmail.com>
---
 dvc/command/data_sync.py  |  6 ++++++
 dvc/repo/add.py           | 13 ++-----------
 dvc/repo/pull.py          | 11 +++++++++--
 dvc/utils/__init__.py     | 13 +++++++++++++
 tests/func/test_import.py | 16 ++++++++++++++++
 5 files changed, 46 insertions(+), 13 deletions(-)

diff --git a/dvc/command/data_sync.py b/dvc/command/data_sync.py
index a7697f152..9fe90dd07 100644
--- a/dvc/command/data_sync.py
+++ b/dvc/command/data_sync.py
@@ -179,6 +179,12 @@ def add_parser(subparsers, _parent_parser):
         default=False,
         help="Fetch run history for all stages.",
     )
+    pull_parser.add_argument(
+        "--glob",
+        action="store_true",
+        default=False,
+        help="Pull cache for targets matching shell-style wildcards.",
+    )
     pull_parser.set_defaults(func=CmdDataPull)
 
     # Push
diff --git a/dvc/repo/add.py b/dvc/repo/add.py
index 47eadbe08..397a33d92 100644
--- a/dvc/repo/add.py
+++ b/dvc/repo/add.py
@@ -14,7 +14,7 @@ from ..exceptions import (
 from ..output.base import OutputDoesNotExistError
 from ..progress import Tqdm
 from ..repo.scm_context import scm_context
-from ..utils import LARGE_DIR_SIZE, resolve_paths
+from ..utils import LARGE_DIR_SIZE, glob_targets, resolve_paths
 from . import locked
 
 logger = logging.getLogger(__name__)
@@ -152,18 +152,9 @@ def _find_all_targets(repo, target, recursive):
 def _create_stages(
     repo, targets, fname, pbar=None, external=False, glob=False, desc=None,
 ):
-    from glob import iglob
-
     from dvc.stage import Stage, create_stage, restore_meta
 
-    if glob:
-        expanded_targets = [
-            exp_target
-            for target in targets
-            for exp_target in iglob(target, recursive=True)
-        ]
-    else:
-        expanded_targets = targets
+    expanded_targets = glob_targets(targets, glob=glob)
 
     stages = []
     for out in Tqdm(
diff --git a/dvc/repo/pull.py b/dvc/repo/pull.py
index a49d16938..e0714da3f 100644
--- a/dvc/repo/pull.py
+++ b/dvc/repo/pull.py
@@ -1,6 +1,7 @@
 import logging
 
 from dvc.repo import locked
+from dvc.utils import glob_targets
 
 logger = logging.getLogger(__name__)
 
@@ -18,12 +19,15 @@ def pull(
     recursive=False,
     all_commits=False,
     run_cache=False,
+    glob=False,
 ):
     if isinstance(targets, str):
         targets = [targets]
 
+    expanded_targets = glob_targets(targets, glob=glob)
+
     processed_files_count = self.fetch(
-        targets,
+        expanded_targets,
         jobs,
         remote=remote,
         all_branches=all_branches,
@@ -34,7 +38,10 @@ def pull(
         run_cache=run_cache,
     )
     stats = self.checkout(
-        targets=targets, with_deps=with_deps, force=force, recursive=recursive
+        targets=expanded_targets,
+        with_deps=with_deps,
+        force=force,
+        recursive=recursive,
     )
 
     stats["fetched"] = processed_files_count
diff --git a/dvc/utils/__init__.py b/dvc/utils/__init__.py
index 993319037..351826756 100644
--- a/dvc/utils/__init__.py
+++ b/dvc/utils/__init__.py
@@ -463,3 +463,16 @@ def parse_target(
 
 def is_exec(mode):
     return mode & (stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)
+
+
+def glob_targets(targets, glob=True, recursive=True):
+    if not glob:
+        return targets
+
+    from glob import iglob
+
+    return [
+        exp_target
+        for target in targets
+        for exp_target in iglob(target, recursive=recursive)
+    ]
diff --git a/tests/func/test_import.py b/tests/func/test_import.py
index a4ad0014e..9d945053f 100644
--- a/tests/func/test_import.py
+++ b/tests/func/test_import.py
@@ -248,6 +248,22 @@ def test_pull_imported_directory_stage(tmp_dir, dvc, erepo_dir):
     assert (tmp_dir / "dir_imported").read_text() == {"foo": "foo content"}
 
 
+def test_pull_wildcard_imported_directory_stage(tmp_dir, dvc, erepo_dir):
+    with erepo_dir.chdir():
+        erepo_dir.dvc_gen(
+            {"dir123": {"foo": "foo content"}}, commit="create dir"
+        )
+
+    dvc.imp(os.fspath(erepo_dir), "dir123", "dir_imported123")
+
+    remove("dir_imported123")
+    remove(dvc.cache.local.cache_dir)
+
+    dvc.pull(["dir_imported*.dvc"], glob=True)
+
+    assert (tmp_dir / "dir_imported123").read_text() == {"foo": "foo content"}
+
+
 def test_download_error_pulling_imported_stage(tmp_dir, dvc, erepo_dir):
     with erepo_dir.chdir():
         erepo_dir.dvc_gen("foo", "foo content", commit="create foo")