From 0f189979467de61cb0512cc0a133d2122e041e90 Mon Sep 17 00:00:00 2001
From: UebelAndre <github@uebelandre.com>
Date: Sun, 20 Mar 2022 12:48:52 -0700
Subject: [PATCH] Added support for custom `strip_prefix` args in toolchains
 (#664)

---
 python/repositories.bzl | 10 ++++++++--
 python/versions.bzl     | 12 ++++++++++--
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/python/repositories.bzl b/python/repositories.bzl
index 02b5011747..d4b113f596 100644
--- a/python/repositories.bzl
+++ b/python/repositories.bzl
@@ -81,7 +81,7 @@ def _python_repository_impl(rctx):
         rctx.download_and_extract(
             url = url,
             sha256 = rctx.attr.sha256,
-            stripPrefix = "python",
+            stripPrefix = rctx.attr.strip_prefix,
         )
 
     # Write distutils.cfg to the Python installation.
@@ -157,6 +157,7 @@ py_runtime_pair(
         "python_version": python_version,
         "release_filename": release_filename,
         "sha256": rctx.attr.sha256,
+        "strip_prefix": rctx.attr.strip_prefix,
         "url": url,
     }
 
@@ -192,6 +193,10 @@ python_repository = repository_rule(
             doc = "The SHA256 integrity hash for the Python interpreter tarball.",
             mandatory = True,
         ),
+        "strip_prefix": attr.string(
+            doc = "A directory prefix to strip from the extracted files.",
+            mandatory = True,
+        ),
         "url": attr.string(
             doc = "The URL of the interpreter to download",
             mandatory = True,
@@ -244,7 +249,7 @@ def python_register_toolchains(
         if not sha256:
             continue
 
-        (release_filename, url) = get_release_url(platform, python_version, base_url, tool_versions)
+        (release_filename, url, strip_prefix) = get_release_url(platform, python_version, base_url, tool_versions)
 
         python_repository(
             name = "{name}_{platform}".format(
@@ -258,6 +263,7 @@ def python_register_toolchains(
             url = url,
             distutils = distutils,
             distutils_content = distutils_content,
+            strip_prefix = strip_prefix,
             **kwargs
         )
         native.register_toolchains("@{name}_toolchains//:{platform}_toolchain".format(
diff --git a/python/versions.bzl b/python/versions.bzl
index c5445ec030..ba1ce5ee36 100644
--- a/python/versions.bzl
+++ b/python/versions.bzl
@@ -34,6 +34,7 @@ TOOL_VERSIONS = {
             "x86_64-apple-darwin": "8d06bec08db8cdd0f64f4f05ee892cf2fcbc58cfb1dd69da2caab78fac420238",
             "x86_64-unknown-linux-gnu": "aec8c4c53373b90be7e2131093caa26063be6d9d826f599c935c0e1042af3355",
         },
+        "strip_prefix": "python",
     },
     "3.8.12": {
         "url": "20220227/cpython-{python_version}+20220227-{platform}-{build}.tar.gz",
@@ -43,6 +44,7 @@ TOOL_VERSIONS = {
             "x86_64-pc-windows-msvc": "924f9fd51ff6ccc533ed8e96c5461768da5781eb3dfc11d846f9e300fab44eda",
             "x86_64-unknown-linux-gnu": "5be9c6d61e238b90dfd94755051c0d3a2d8023ebffdb4b0fa4e8fedd09a6cab6",
         },
+        "strip_prefix": "python",
     },
     "3.9.10": {
         "url": "20220227/cpython-{python_version}+20220227-{platform}-{build}.tar.gz",
@@ -52,6 +54,7 @@ TOOL_VERSIONS = {
             "x86_64-pc-windows-msvc": "5bc65ce023614bf496a6748e41dca934b70fc5fac6dfacc46aa8dbcad772afc2",
             "x86_64-unknown-linux-gnu": "455089cc576bd9a58db45e919d1fc867ecdbb0208067dffc845cc9bbf0701b70",
         },
+        "strip_prefix": "python",
     },
     "3.10.2": {
         "url": "20220227/cpython-{python_version}+20220227-{platform}-{build}.tar.gz",
@@ -61,6 +64,7 @@ TOOL_VERSIONS = {
             "x86_64-pc-windows-msvc": "a293c5838dd9c8438a84372fb95dda9752df63928a8a2ae516438f187f89567d",
             "x86_64-unknown-linux-gnu": "9b64eca2a94f7aff9409ad70bdaa7fbbf8148692662e764401883957943620dd",
         },
+        "strip_prefix": "python",
     },
 }
 
@@ -118,7 +122,7 @@ def get_release_url(platform, python_version, base_url = DEFAULT_RELEASE_BASE_UR
         tool_versions: A dict listing the interpreter versions, their SHAs and URL
 
     Returns:
-        filename and url pair
+        A tuple of (filename, url, and archive strip prefix)
     """
 
     url = tool_versions[python_version]["url"]
@@ -126,13 +130,17 @@ def get_release_url(platform, python_version, base_url = DEFAULT_RELEASE_BASE_UR
     if type(url) == type({}):
         url = url[platform]
 
+    strip_prefix = tool_versions[python_version].get("strip_prefix", None)
+    if type(strip_prefix) == type({}):
+        strip_prefix = strip_prefix[platform]
+
     release_filename = url.format(
         platform = platform,
         python_version = python_version,
         build = "static-install_only" if (WINDOWS_NAME in platform) else "install_only",
     )
     url = "/".join([base_url, release_filename])
-    return (release_filename, url)
+    return (release_filename, url, strip_prefix)
 
 def print_toolchains_checksums(name):
     native.genrule(