diff --git a/dvc/commands/data_sync.py b/dvc/commands/data_sync.py index aab2a6459..117684afd 100644 --- a/dvc/commands/data_sync.py +++ b/dvc/commands/data_sync.py @@ -91,6 +91,7 @@ class CmdDataFetch(CmdDataBase): with_deps=self.args.with_deps, recursive=self.args.recursive, run_cache=self.args.run_cache, + max_size=self.args.max_size, ) self.log_summary({"fetched": processed_files_count}) except DvcException: @@ -322,6 +323,11 @@ def add_parser(subparsers, _parent_parser): default=False, help="Fetch run history for all stages.", ) + fetch_parser.add_argument( + "--max-size", + type=int, + help="Fetch data files/directories that are each below specified size (bytes).", + ) fetch_parser.set_defaults(func=CmdDataFetch) # Status diff --git a/dvc/repo/fetch.py b/dvc/repo/fetch.py index 7a061c5cc..5df04e6de 100644 --- a/dvc/repo/fetch.py +++ b/dvc/repo/fetch.py @@ -22,6 +22,7 @@ def fetch( # noqa: C901, PLR0913 all_commits=False, run_cache=False, revs=None, + max_size=None, ) -> int: """Download data items from a cloud and imported repositories @@ -70,6 +71,7 @@ def fetch( # noqa: C901, PLR0913 targets, with_deps=with_deps, recursive=recursive, + max_size=max_size, ) index_keys.add(idx.data_tree.hash_info.value) indexes.append(idx.data["repo"]) diff --git a/dvc/repo/index.py b/dvc/repo/index.py index da8c03bf0..bb9a704ee 100644 --- a/dvc/repo/index.py +++ b/dvc/repo/index.py @@ -492,6 +492,7 @@ class Index: targets: Optional["TargetType"], stage_filter: Optional[Callable[["Stage"], bool]] = None, outs_filter: Optional[Callable[["Output"], bool]] = None, + max_size: Optional[int] = None, **kwargs: Any, ) -> "IndexView": """Return read-only view of index for the specified targets. @@ -514,7 +515,17 @@ class Index: for stage_info in self.collect_targets(targets, **kwargs) if not stage_filter or stage_filter(stage_info.stage) ] - return IndexView(self, stage_infos, outs_filter=outs_filter) + + def _outs_filter(out): + if max_size and out.meta and out.meta.size and out.meta.size >= max_size: + return False + + if outs_filter: + return outs_filter(out) + + return True + + return IndexView(self, stage_infos, outs_filter=_outs_filter) class _DataPrefixes(NamedTuple): diff --git a/tests/unit/command/test_data_sync.py b/tests/unit/command/test_data_sync.py index e5f439b6d..9af90c529 100644 --- a/tests/unit/command/test_data_sync.py +++ b/tests/unit/command/test_data_sync.py @@ -18,6 +18,8 @@ def test_fetch(mocker, dvc): "--with-deps", "--recursive", "--run-cache", + "--max-size", + "10", ] ) assert cli_args.func == CmdDataFetch @@ -37,6 +39,7 @@ def test_fetch(mocker, dvc): with_deps=True, recursive=True, run_cache=True, + max_size=10, )