fetch: introduce --max-size (#9696)

Needed for iterative/studio#4782

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

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"])

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):

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,
     )