From 3e1814e5ec13590b188c9d805bd150b5ab653ebf Mon Sep 17 00:00:00 2001
From: Matteo De Wint <matteo@gorilla.co>
Date: Sun, 31 Dec 2023 14:07:54 +0100
Subject: [PATCH] refactor: don't walk directories recursively

---
 s3pypi/storage.py                 | 25 ++++++++-----------------
 tests/integration/test_storage.py |  8 +-------
 2 files changed, 9 insertions(+), 24 deletions(-)

diff --git a/s3pypi/storage.py b/s3pypi/storage.py
index 65a0b50..35885c6 100644
--- a/s3pypi/storage.py
+++ b/s3pypi/storage.py
@@ -1,4 +1,3 @@
-from collections import deque
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Dict, List, Optional
@@ -53,22 +52,14 @@ def build_root_index(self) -> Index:
         return Index(dict.fromkeys(self._list_dirs()))
 
     def _list_dirs(self) -> List[str]:
-        results = set()
-        root = f"{p}/" if (p := self.cfg.prefix) else ""
-        todo = deque([root])
-        while todo:
-            current = todo.popleft()
-            if children := [
-                prefix
-                for item in self.s3.meta.client.get_paginator("list_objects_v2")
-                .paginate(Bucket=self.cfg.bucket, Delimiter="/", Prefix=current)
-                .search("CommonPrefixes")
-                if item and (prefix := item.get("Prefix"))
-            ]:
-                todo.extend(children)
-            else:
-                results.add(current[len(root) :])
-        return sorted(results)
+        prefix = f"{p}/" if (p := self.cfg.prefix) else ""
+        return [
+            d[len(prefix) :]
+            for item in self.s3.meta.client.get_paginator("list_objects_v2")
+            .paginate(Bucket=self.cfg.bucket, Delimiter="/", Prefix=prefix)
+            .search("CommonPrefixes")
+            if item and (d := item.get("Prefix"))
+        ]
 
     def put_index(self, directory: str, index: Index) -> None:
         self._object(directory, self.index_name).put(
diff --git a/tests/integration/test_storage.py b/tests/integration/test_storage.py
index 11c60cd..1550f57 100644
--- a/tests/integration/test_storage.py
+++ b/tests/integration/test_storage.py
@@ -60,10 +60,4 @@ def test_list_dirs(boto3_session, s3_bucket):
     cfg = S3Config(bucket=s3_bucket.name)
     s = S3Storage(boto3_session, cfg)
 
-    assert s._list_dirs() == [
-        "AA/one/",
-        "AA/three/",
-        "AA/two/",
-        "BBBB/xxx/",
-        "BBBB/yyy/",
-    ]
+    assert s._list_dirs() == ["AA/", "BBBB/"]