Support repodata v3 draft for monolithic, sharded repodata including .whl

conda_index/api.py

@@ -17,6 +17,7 @@ def update_index(
     progress=False,
     current_index_versions=None,
     write_run_exports=False,
+    repodata_v3=False,
 ):
     import os
 
@@ -49,4 +50,5 @@ def update_index(
             subdirs=ensure_list(subdir),
             current_index_versions=current_index_versions,
             write_run_exports=write_run_exports,
+            repodata_v3=repodata_v3,
         )

conda_index/cli/__init__.py

@@ -185,6 +185,15 @@
     default=False,
     show_default=True,
 )
+@click.option(
+    "--repodata-next/--no-repodata-next",
+    help="""
+        EXPERIMENTAL. Write CEP-XXXX v3 repodata layout with all package
+        records under a top-level v3 key.
+        """,
+    default=False,
+    show_default=True,
+)
 def cli(
     dir,
     patch_generator=None,
@@ -210,6 +219,7 @@ def cli(
     db_url="",
     html_dependencies=False,
     update_only=False,
+    repodata_next=False,
 ):
     logutil.configure()
     if verbose:
@@ -269,6 +279,7 @@ def cli(
         cache_kwargs=cache_kwargs,
         html_dependencies=html_dependencies,
         update_only=update_only,
+        repodata_v3=repodata_next,
     )
 
     if update_cache is False:

conda_index/index/__init__.py

@@ -12,11 +12,10 @@
 import sys
 import time
 from concurrent.futures import Executor, ProcessPoolExecutor, ThreadPoolExecutor
-from contextlib import nullcontext
 from datetime import datetime, timezone
 from os.path import basename, getmtime, getsize, isfile, join
 from pathlib import Path
-from typing import Iterable
+from typing import TYPE_CHECKING, Iterable
 from uuid import uuid4
 
 import msgpack
@@ -36,6 +35,27 @@
 from . import rss, sqlitecache
 from .fs import FileInfo, MinimalFS
 
+if TYPE_CHECKING:
+    from typing import Any, NotRequired, TypedDict
+
+    from .cache import IndexedPackages, IndexedShard
+
+    V3Section = TypedDict(
+        "V3Section",
+        {"tar.bz2": dict, "conda": dict, "whl": dict},
+    )
+
+    # in this style because "packages.conda" is not a Python identifier
+    ShardDict = TypedDict(
+        "ShardDict",
+        {
+            "packages": dict[str, dict[str, Any]],
+            "packages.conda": dict[str, dict[str, Any]],
+            "v3": NotRequired[V3Section],
+        },
+    )
+
+
 log = logging.getLogger(__name__)
 
 # zstd -T0 -b15 -e17 repodata.json
@@ -111,6 +131,7 @@ def update_index(
     write_zst=False,
     write_run_exports=False,
     html_dependencies=False,
+    repodata_v3=False,
 ):
     """
     High-level interface to ``ChannelIndex``. Index all subdirs under
@@ -147,6 +168,7 @@ def update_index(
         write_zst=write_zst,
         write_run_exports=write_run_exports,
         html_dependencies=html_dependencies,
+        repodata_v3=repodata_v3,
     )
 
     channel_index.index(
@@ -177,6 +199,7 @@ def _make_seconds(timestamp):
 )
 CHANNELDATA_VERSION = 1
 RUN_EXPORTS_VERSION = 1
+REPODATA_REVISION_V3 = 3
 REPODATA_JSON_FN = "repodata.json"
 REPODATA_FROM_PKGS_JSON_FN = "repodata_from_packages.json"
 REPODATA_SHARDS_FN = "repodata_shards.msgpack.zst"
@@ -252,15 +275,17 @@ def _apply_instructions(subdir, repodata, instructions, new_pkg_fixes=None):
         for key in ("packages", "packages.conda"):
             if key == "packages.conda" and fn.endswith(CONDA_PACKAGE_EXTENSION_V1):
                 fn = fn.replace(CONDA_PACKAGE_EXTENSION_V1, CONDA_PACKAGE_EXTENSION_V2)
-            if fn in repodata[key]:
-                repodata[key][fn]["revoked"] = True
-                repodata[key][fn]["depends"].append("package_has_been_revoked")
+            records = repodata.get(key, {})
+            if fn in records:
+                records[fn]["revoked"] = True
+                records[fn]["depends"].append("package_has_been_revoked")
 
     for fn in instructions.get("remove", ()):
         for key in ("packages", "packages.conda"):
             if key == "packages.conda" and fn.endswith(CONDA_PACKAGE_EXTENSION_V1):
                 fn = fn.replace(CONDA_PACKAGE_EXTENSION_V1, CONDA_PACKAGE_EXTENSION_V2)
-            popped = repodata[key].pop(fn, None)
+            records = repodata.get(key, {})
+            popped = records.pop(fn, None)
             if popped:
                 repodata["removed"].append(fn)
     repodata["removed"].sort()
@@ -397,6 +422,7 @@ def __init__(
         upstream_stage: str = "fs",
         cache_kwargs=None,
         update_only=False,
+        repodata_v3=False,
     ):
         if threads is None:
             threads = MAX_THREADS_DEFAULT
@@ -430,6 +456,7 @@ def __init__(
         self.write_current_repodata = write_current_repodata
         self.upstream_stage = upstream_stage
         self.update_only = update_only
+        self.repodata_v3 = repodata_v3
 
         self.cache_kwargs = cache_kwargs
 
@@ -717,13 +744,29 @@ def index_subdir_shards(self, subdir, verbose=False, progress=False):
 
         (self.output_root / subdir).mkdir(parents=True, exist_ok=True)
 
-        for name, shard in cache.indexed_shards():
-            shard_data = compressor.compress(sqlitecache.packb_typed(shard))
-            shard_hash = hashlib.sha256(shard_data).digest()
+        v3_data = {
+            "tar.bz2": {},
+            "conda": {},
+            "whl": {},
+        }
+
+        for shard in cache.indexed_shards_2():
+            repodata_shard = self._indexed_shard_to_repodata(shard)
+            shard_bytes = compressor.compress(sqlitecache.packb_typed(repodata_shard))
+            shard_hash = hashlib.sha256(shard_bytes).digest()
             output_path = self.output_root / subdir / f"{shard_hash.hex()}.msgpack.zst"
             if not output_path.exists():
-                output_path.write_bytes(shard_data)
-            shards[name] = shard_hash
+                output_path.write_bytes(shard_bytes)
+            shards[shard.name] = shard_hash
+
+            if self.repodata_v3:
+                for section, records in repodata_shard["v3"].items():
+                    v3_data[section].update(records)
+
+        if self.repodata_v3:
+            shards_index["info"]["repodata_revisions"] = [
+                self._make_repodata_revision_data(v3_data)
+            ]
 
         return shards_index
 
@@ -750,13 +793,108 @@ def index_subdir(self, subdir, verbose=False, progress=False):
             "removed": [],  # can be added by patch/hotfix process
         }
 
+        if self.repodata_v3:
+            v3_packages = self._extract_indexed_packages_v3(indexed_packages)
+            new_repodata["v3"] = v3_packages
+            new_repodata["packages"] = {}
+            new_repodata["packages.conda"] = {}
+            new_repodata["info"]["repodata_revisions"] = [
+                self._make_repodata_revision_data(v3_packages)
+            ]
+
         if self.base_url:
             # per https://github.com/conda-incubator/ceps/blob/main/cep-15.md
             new_repodata["info"]["base_url"] = f"{self.base_url.rstrip('/')}/{subdir}/"
             new_repodata["repodata_version"] = 2
 
         return new_repodata
 
+    @staticmethod
+    def _v3_key_for_path(path: str) -> str | None:
+        for extension in (".tar.bz2", ".conda", ".whl"):
+            if path.endswith(extension):
+                return path[: -len(extension)]
+        return None
+
+    def _extract_indexed_packages_v3(
+        self, indexed_packages: IndexedPackages
+    ) -> V3Section:
+        """
+        Return all packages from IndexedPackages as the "v3": {...} section.
+        """
+        v3: V3Section = {
+            "tar.bz2": {},
+            "conda": {},
+            "whl": {},
+        }
+        for section, records in (
+            ("tar.bz2", indexed_packages.packages),
+            ("conda", indexed_packages.packages_conda),
+            ("whl", indexed_packages.packages_whl),
+        ):
+            for filename, record in records.items():
+                # Per draft wheel-in-conda work, key is conda-like so that some
+                # conda-like parsing can occur on the key only. So we derive the
+                # key here. `record["fn"]` contains the filename or URL.
+                if section == "whl":
+                    name = record.get("name")
+                    version = record.get("version")
+                    build = record.get("build")
+                    if name is None or version is None or build is None:
+                        log.warning(
+                            "%s: v3 whl records require name, version, and build; skipping",
+                            filename,
+                        )
+                        continue
+                    key = f"{name}-{version}-{build}"
+                else:
+                    key = self._v3_key_for_path(filename)
+                    if key is None:
+                        log.warning("%s has unsupported package extension", filename)
+                        continue
+
+                v3[section][key] = record
+
+        return v3
+
+    def _indexed_shard_to_repodata(self, indexed_shard: IndexedShard) -> ShardDict:
+        if self.repodata_v3:
+            shard_data: ShardDict = {
+                "packages": {},
+                "packages.conda": {},
+                "v3": self._extract_indexed_packages_v3(indexed_shard),
+            }
+        else:
+            shard_data = {
+                "packages": indexed_shard.packages,
+                "packages.conda": indexed_shard.packages_conda,
+            }
+        return shard_data
+
+    @staticmethod
+    def _make_repodata_revision_data(
+        revision_data: dict[str, dict[str, dict]],
+    ) -> dict[str, int | None]:
+        """
+        Return { "revision": 3, ... } dict with package statistics derived from
+        revision_data, which is similar to monolithic repodata.
+        """
+        timestamps = []
+        n_packages = 0
+        for section_records in revision_data.values():
+            n_packages += len(section_records)
+            for record in section_records.values():
+                timestamp = record.get("timestamp")
+                if isinstance(timestamp, (int, float)):
+                    timestamps.append(int(timestamp))
+
+        return {
+            "revision": REPODATA_REVISION_V3,
+            "n_packages": n_packages,
+            "oldest": min(timestamps) if timestamps else None,
+            "newest": max(timestamps) if timestamps else None,
+        }
+
     def extract_subdir_to_cache(
         self,
         subdir: str,
@@ -826,7 +964,7 @@ def extract_subdir_to_cache(
 
         return subdir
 
-    ####
+    # region: channeldata
 
     def channeldata_path(self):
         channeldata_file = os.path.join(self.output_root, "channeldata.json")
@@ -869,6 +1007,8 @@ def update_channeldata(self, rss=False):
         log.debug("write channeldata")
         self._write_channeldata(channel_data)
 
+    # endregion
+
     def detect_subdirs(self):
         if not self._subdirs:
             detected_subdirs = {