Add progress

src/pypi_data/cli.py

@@ -81,30 +81,30 @@ def load_repos(
         repos = repos[:limit]
 
     repos = asyncio.run(load_indexes(repos))
-    sorted_repos = sorted(repos, key=lambda r: r.number)
+    repos = sorted(repos, key=lambda r: r.number)
 
     log.info("Loaded: writing file")
     with open_path(repos_file, mode="wb") as fd:
-        for repo in repos:
+        for repo in tqdm.tqdm(repos):
             fd.write(repo.model_dump_json().encode("utf-8"))
             fd.write(b"\n")
 
     log.info("Writing links")
 
     (links_path / "dataset.txt").write_text(
-        "\n".join(str(repo.dataset_url) for repo in sorted_repos)
+        "\n".join(str(repo.dataset_url) for repo in repos)
     )
 
     (links_path / "repositories.txt").write_text(
-        "\n".join(str(repo.url) for repo in sorted_repos)
+        "\n".join(str(repo.url) for repo in repos)
     )
 
     (links_path / "repositories_ssh.txt").write_text(
-        "\n".join(str(repo.ssh_url) for repo in sorted_repos)
+        "\n".join(str(repo.ssh_url) for repo in repos)
     )
 
     (links_path / "repositories.json").write_text(
-        Repos([r.without_index() for r in sorted_repos]).model_dump_json(
+        Repos([r.without_index() for r in repos]).model_dump_json(
             indent=2, exclude_none=True
         )
     )
@@ -113,7 +113,7 @@ def load_repos(
 
     package_map: DefaultDict[str, list[PackageIndexPackage]] = defaultdict(list)
 
-    for repo in sorted_repos:
+    for repo in repos:
         for package in repo.index.packages:
             package_map[package.project_name].append(
                 PackageIndexPackage(
@@ -135,7 +135,7 @@ def load_repos(
     log.info("Writing package index")
 
     with open_path(packages_file, mode="wb") as fd:
-        for package_index in packages:
+        for package_index in tqdm.tqdm(packages):
             fd.write(package_index.model_dump_json().encode("utf-8"))
             fd.write(b"\n")
 

src/pypi_data/combine_parquet.py

@@ -1,5 +1,7 @@
 import hashlib
+from collections import deque
 from pathlib import Path
+from typing import Deque
 
 import httpx
 import pyarrow
@@ -30,15 +32,15 @@ def append_buffer(
 
 
 async def fill_buffer(
-    buffer: list[tuple[tuple[int, str], RecordBatch]],
+    buffer: Deque[tuple[tuple[int, str], RecordBatch]],
     client: httpx.AsyncClient,
-    repositories: list[CodeRepository],
+    repositories: Deque[CodeRepository],
     path: Path,
 ) -> bool:
     for _ in range(FILL_BUFFER_COUNT):
         if not repositories:
             break
-        repo = repositories.pop(0)
+        repo = repositories.popleft()
         log.info(f"Downloading {repo.dataset_url}")
         if await repo.download_dataset(client, path) is False:
             log.info(f"Failed to download {repo.dataset_url}")
@@ -67,7 +69,8 @@ async def combine_parquet(repositories: list[CodeRepository], directory: Path):
     repo_file = directory / "repo.parquet"
 
     roll_up_count = 0
-    buffer: list[tuple[tuple[int, str], RecordBatch]] = []
+    buffer: Deque[tuple[tuple[int, str], RecordBatch]] = deque()
+    repositories: Deque[CodeRepository] = deque(repositories)
 
     async with httpx.AsyncClient(follow_redirects=True) as client:
         while repositories:
@@ -77,7 +80,7 @@ async def combine_parquet(repositories: list[CodeRepository], directory: Path):
             roll_up_path = directory / f"merged-{roll_up_count}.parquet"
 
             keys = []
-            first_key, first_buffer = buffer.pop(0)
+            first_key, first_buffer = buffer.popleft()
             keys.append(first_key)
 
             with pq.ParquetWriter(
@@ -98,7 +101,7 @@ async def combine_parquet(repositories: list[CodeRepository], directory: Path):
                         ):
                             continue
 
-                    key, batch = buffer.pop(0)
+                    key, batch = buffer.popleft()
                     keys.append(key)
                     if append_buffer(writer, batch, roll_up_path):
                         break