refactor(pypi): factor out a simple implementation of the PyPI cache (#3639)

We want to keep a `dict` like interface and later we would like to use
the same interface to also do more things. I expect the cache key to
become different in the future (i.e. include requested versions in it)
so that we can check if we have the right versions in the
MODULE.bazel.lock file or if we should actually call to PyPI.

Work towards #2731

Author: aignas

python/private/pypi/BUILD.bazel

@@ -123,6 +123,7 @@ bzl_library(
         ":pep508_env_bzl",
         ":pip_repository_attrs_bzl",
         ":platform_bzl",
+        ":pypi_cache_bzl",
         ":simpleapi_download_bzl",
         ":whl_library_bzl",
         "//python/private:auth_bzl",
@@ -355,6 +356,11 @@ bzl_library(
     srcs = ["platform.bzl"],
 )
 
+bzl_library(
+    name = "pypi_cache_bzl",
+    srcs = ["pypi_cache.bzl"],
+)
+
 bzl_library(
     name = "pypi_repo_utils_bzl",
     srcs = ["pypi_repo_utils.bzl"],

python/private/pypi/extension.bzl

@@ -27,6 +27,7 @@ load(":parse_whl_name.bzl", "parse_whl_name")
 load(":pep508_env.bzl", "env")
 load(":pip_repository_attrs.bzl", "ATTRS")
 load(":platform.bzl", _plat = "platform")
+load(":pypi_cache.bzl", "pypi_cache")
 load(":simpleapi_download.bzl", "simpleapi_download")
 load(":whl_library.bzl", "whl_library")
 
@@ -224,7 +225,7 @@ You cannot use both the additive_build_content and additive_build_content_file a
     # dict[str repo, HubBuilder]
     # See `hub_builder.bzl%hub_builder()` for `HubBuilder`
     pip_hub_map = {}
-    simpleapi_cache = {}
+    simpleapi_cache = pypi_cache()
 
     for mod in module_ctx.modules:
         for pip_attr in mod.tags.parse:

python/private/pypi/hub_builder.bzl

@@ -31,7 +31,7 @@ def hub_builder(
         simpleapi_download_fn,
         evaluate_markers_fn,
         logger,
-        simpleapi_cache = {}):
+        simpleapi_cache):
     """Return a hub builder instance
 
     Args:

python/private/pypi/pypi_cache.bzl

@@ -0,0 +1,47 @@
+"""A cache for the PyPI index contents evaluation.
+
+This is design to work as the following:
+- in-memory cache for results of PyPI index queries, so that we are not calling PyPI multiple times
+  for the same package for different hub repos.
+
+In the future the same will be used to:
+- Store PyPI index query results as facts in the MODULE.bazel.lock file
+"""
+
+def pypi_cache(store = None):
+    """The cache for PyPI index queries."""
+
+    # buildifier: disable=uninitialized
+    self = struct(
+        _store = store or {},
+        setdefault = lambda key, parsed_result: _pypi_cache_setdefault(self, key, parsed_result),
+        get = lambda key: _pypi_cache_get(self, key),
+    )
+
+    # buildifier: enable=uninitialized
+    return self
+
+def _pypi_cache_setdefault(self, key, parsed_result):
+    """Store the value if not yet cached.
+
+    Args:
+        self: {type}`struct` The self of this implementation.
+        key: {type}`str` The cache key, can be any string.
+        parsed_result: {type}`struct` The result of `parse_simpleapi_html` function.
+
+    Returns:
+        The `parse_result`.
+    """
+    return self._store.setdefault(key, parsed_result)
+
+def _pypi_cache_get(self, key):
+    """Return the parsed result from the cache.
+
+    Args:
+        self: {type}`struct` The self of this implementation.
+        key: {type}`str` The cache key, can be any string.
+
+    Returns:
+        The {type}`struct` or `None` based on if the result is in the cache or not.
+    """
+    return self._store.get(key)

python/private/pypi/simpleapi_download.bzl

@@ -49,14 +49,13 @@ def simpleapi_download(
            * netrc: The netrc parameter for ctx.download, see http_file for docs.
            * auth_patterns: The auth_patterns parameter for ctx.download, see
                http_file for docs.
-        cache: A dictionary that can be used as a cache between calls during a
-            single evaluation of the extension. We use a dictionary as a cache
-            so that we can reuse calls to the simple API when evaluating the
-            extension. Using the canonical_id parameter of the module_ctx would
-            deposit the simple API responses to the bazel cache and that is
-            undesirable because additions to the PyPI index would not be
-            reflected when re-evaluating the extension unless we do
-            `bazel clean --expunge`.
+        cache: An opaque object used to cache call results. For implementation
+            see ./pypi_cache.bzl file. We use the canonical_id parameter for the key
+            value to ensure that distribution fetches from different indexes do not cause
+            cache collisions, because the index may return different locations from where
+            the files should be downloaded. We are not using the built-in cache in the
+            `download` function because the index may get updated at any time and we need
+            to be able to refresh the data.
         parallel_download: A boolean to enable usage of bazel 7.1 non-blocking downloads.
         read_simpleapi: a function for reading and parsing of the SimpleAPI contents.
             Used in tests.
@@ -197,8 +196,9 @@ def _read_simpleapi(ctx, url, attr, cache, get_auth = None, **download_kwargs):
     ))
 
     cache_key = real_url
-    if cache_key in cache:
-        return struct(success = True, output = cache[cache_key])
+    cached_result = cache.get(cache_key)
+    if cached_result:
+        return struct(success = True, output = cached_result)
 
     output_str = envsubst(
         url,

tests/pypi/hub_builder/hub_builder_tests.bzl

@@ -99,6 +99,7 @@ def hub_builder(
             "unit-test",
             printer = log_printer,
         ),
+        simpleapi_cache = {},
     )
     self = struct(
         build = lambda: env.expect.that_struct(

tests/pypi/simpleapi_download/simpleapi_download_tests.bzl

@@ -15,6 +15,7 @@
 ""
 
 load("@rules_testing//lib:test_suite.bzl", "test_suite")
+load("//python/private/pypi:pypi_cache.bzl", "pypi_cache")  # buildifier: disable=bzl-visibility
 load("//python/private/pypi:simpleapi_download.bzl", "simpleapi_download", "strip_empty_path_segments")  # buildifier: disable=bzl-visibility
 
 _tests = []
@@ -52,7 +53,7 @@ def _test_simple(env):
             sources = ["foo", "bar", "baz"],
             envsubst = [],
         ),
-        cache = {},
+        cache = pypi_cache(),
         parallel_download = True,
         read_simpleapi = read_simpleapi,
     )
@@ -112,7 +113,7 @@ def _test_fail(env):
             sources = ["foo", "bar", "baz"],
             envsubst = [],
         ),
-        cache = {},
+        cache = pypi_cache(),
         parallel_download = True,
         read_simpleapi = read_simpleapi,
         _fail = fails.append,
@@ -165,7 +166,7 @@ def _test_download_url(env):
             sources = ["foo", "bar", "baz"],
             envsubst = [],
         ),
-        cache = {},
+        cache = pypi_cache(),
         parallel_download = False,
         get_auth = lambda ctx, urls, ctx_attr: struct(),
     )
@@ -201,7 +202,7 @@ def _test_download_url_parallel(env):
             sources = ["foo", "bar", "baz"],
             envsubst = [],
         ),
-        cache = {},
+        cache = pypi_cache(),
         parallel_download = True,
         get_auth = lambda ctx, urls, ctx_attr: struct(),
     )
@@ -237,7 +238,7 @@ def _test_download_envsubst_url(env):
             sources = ["foo", "bar", "baz"],
             envsubst = ["INDEX_URL"],
         ),
-        cache = {},
+        cache = pypi_cache(),
         parallel_download = False,
         get_auth = lambda ctx, urls, ctx_attr: struct(),
     )