Add --put-root-index option

CHANGELOG.md

@@ -6,6 +6,14 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [PEP 440](https://www.python.org/dev/peps/pep-0440/).
 
 
+## Unreleased
+
+### Added
+
+- `--lock-indexes` option to lock index objects in S3 using a DynamoDB table.
+- `--put-root-index` option to write a root index that lists all package names.
+
+
 ## 1.0.0rc1 - 2021-05-19
 
 ### Added

s3pypi/__main__.py

@@ -56,6 +56,7 @@ def get_arg_parser():
             "This ensures that concurrent invocations of s3pypi do not overwrite each other's changes."
         ),
     )
+    p.add_argument("--put-root-index", action="store_true", help="Write a root index.")
     p.add_argument("-f", "--force", action="store_true", help="Overwrite files.")
     p.add_argument("-v", "--verbose", action="store_true", help="Verbose output.")
     p.add_argument("-V", "--version", action="version", version=__version__)

s3pypi/core.py

@@ -36,6 +36,7 @@ def upload_packages(
     bucket: str,
     force: bool = False,
     lock_indexes: bool = False,
+    put_root_index: bool = False,
     profile: Optional[str] = None,
     region: Optional[str] = None,
     **kwargs,
@@ -67,11 +68,10 @@ def upload_packages(
 
             storage.put_index(directory, index)
 
-    root = storage.root
-    with lock(root):
-        index = storage.get_index(root)
-        # TODO: Update root index
-        storage.put_index(root, index)
+    if put_root_index:
+        with lock(storage.root):
+            index = storage.build_root_index()
+            storage.put_index(storage.root, index)
 
 
 def parse_distribution(path: Path) -> Distribution:

s3pypi/index.py

@@ -16,7 +16,7 @@ def parse(cls, html: str) -> "Index":
 
     def to_html(self) -> str:
         links = "<br>\n".join(
-            f'<a href="{urllib.parse.quote(fname)}">{fname}</a>'
+            f'<a href="{urllib.parse.quote(fname)}">{fname.rstrip("/")}</a>'
             for fname in sorted(self.filenames)
         )
         return index_html.format(body=indent(links, " " * 4))

s3pypi/storage.py

@@ -44,6 +44,17 @@ def get_index(self, directory: str) -> Index:
             return Index()
         return Index.parse(html.decode())
 
+    def build_root_index(self) -> Index:
+        paginator = self.s3.meta.client.get_paginator("list_objects_v2")
+        result = paginator.paginate(
+            Bucket=self.bucket,
+            Prefix=self.prefix or "",
+            Delimiter="/",
+        )
+        n = len(self.prefix) + 1 if self.prefix else 0
+        dirs = set(p.get("Prefix")[n:] for p in result.search("CommonPrefixes"))
+        return Index(dirs)
+
     def put_index(self, directory: str, index: Index):
         self._object(directory, self.index_name).put(
             Body=index.to_html(),

tests/integration/test_main.py

@@ -23,14 +23,17 @@ def test_string_dict(text, expected):
 def test_main_upload_package(chdir, data_dir, s3_bucket):
     with chdir(data_dir):
         dist = sorted(glob.glob("dists/*"))
-        s3pypi(*dist, "--bucket", s3_bucket.name)
+        s3pypi(*dist, "--bucket", s3_bucket.name, "--put-root-index")
 
     def read(key: str) -> bytes:
         return s3_bucket.Object(key).get()["Body"].read()
 
+    root_index = read("index.html").decode()
+
     def assert_pkg_exists(prefix: str, filename: str):
         assert read(prefix + filename)
         assert f">{filename}</a>" in read(prefix).decode()
+        assert f">{prefix.rstrip('/')}</a>" in root_index
 
     assert_pkg_exists("foo/", "foo-0.1.0.tar.gz")
     assert_pkg_exists("hello-world/", "hello_world-0.1.0-py3-none-any.whl")