From 7fd191126ea2f8c7a11ead921ab33d1e036704a5 Mon Sep 17 00:00:00 2001 From: Batuhan Taskaya Date: Mon, 11 Jan 2021 12:54:42 +0300 Subject: [PATCH] S3Tree._list_paths now recognizes the max_items limit (#5239) * S3Tree._list_paths now recognizes the max_items limit * Use modern approach + add a test * refactor tests Co-authored-by: Ruslan Kuprieiev * lint Co-authored-by: Ruslan Kuprieiev --- dvc/tree/s3.py | 11 ++++++----- tests/func/test_s3.py | 8 ++++++++ 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/dvc/tree/s3.py b/dvc/tree/s3.py index 4d1acb541..ca4efdd4d 100644 --- a/dvc/tree/s3.py +++ b/dvc/tree/s3.py @@ -202,12 +202,13 @@ class S3Tree(BaseTree): return path_info.path in self._list_paths(path_info) def _list_paths(self, path_info, max_items=None): - kwargs = {"Prefix": path_info.path} - if max_items is not None: - kwargs["MaxKeys"] = max_items - with self._get_bucket(path_info.bucket) as bucket: - for obj_summary in bucket.objects.filter(**kwargs): + obj_summaries = bucket.objects.filter(Prefix=path_info.path) + if max_items is not None: + obj_summaries = obj_summaries.page_size(max_items).limit( + max_items + ) + for obj_summary in obj_summaries: yield obj_summary.key def walk_files(self, path_info, **kwargs): diff --git a/tests/func/test_s3.py b/tests/func/test_s3.py index fcc7cc831..9e52aecd3 100644 --- a/tests/func/test_s3.py +++ b/tests/func/test_s3.py @@ -111,3 +111,11 @@ def test_copy_multipart_preserve_etag(): s3.create_bucket(Bucket=from_info.bucket) _upload_multipart(s3, from_info.bucket, from_info.path) S3Tree._copy(s3, from_info, to_info, {}) + + +def test_s3_isdir(tmp_dir, dvc, s3): + s3.gen({"data": {"foo": "foo"}}) + tree = S3Tree(dvc, s3.config) + + assert not tree.isdir(s3 / "data" / "foo") + assert tree.isdir(s3 / "data")