From 07ae4e93094c6c07243403a6b8411cd090c2766d Mon Sep 17 00:00:00 2001
From: Eli Uriegas <eliuriegas@fb.com>
Date: Mon, 1 Mar 2021 16:43:08 -0800
Subject: [PATCH] scripts: Add script to prep wheels for pypi (#53056)

Summary:
Adds a script so that we can take wheels directly from
download.pytorch.org and publish them to pypi

This is currently mainly used to prep windows binaries for publication to PyPI

Signed-off-by: Eli Uriegas <eliuriegas@fb.com>

Pull Request resolved: https://github.com/pytorch/pytorch/pull/53056

Reviewed By: H-Huang

Differential Revision: D26738642

Pulled By: seemethere

fbshipit-source-id: 96777ed6c3f3454bddb4bc13121f727074312816
---
 .../release/promote/prep_binary_for_pypi.sh   | 43 +++++++++++++++++++
 1 file changed, 43 insertions(+)
 create mode 100755 scripts/release/promote/prep_binary_for_pypi.sh

diff --git a/scripts/release/promote/prep_binary_for_pypi.sh b/scripts/release/promote/prep_binary_for_pypi.sh
new file mode 100755
index 00000000000..041fb109e32
--- /dev/null
+++ b/scripts/release/promote/prep_binary_for_pypi.sh
@@ -0,0 +1,43 @@
+#!/usr/bin/env bash
+
+# Preps binaries for publishing to pypi by removing the
+# version suffix we normally add for all binaries
+# (outside of default ones, CUDA 10.2 currently)
+
+# Usage is:
+# $ prep_binary_for_pypy.sh <path_to_whl_file> <path_to_multiple_whl_files>
+
+# Will output a whl in your current directory
+
+set -eou pipefail
+shopt -s globstar
+
+tmp_dir="$(mktemp -d)"
+trap 'rm -rf ${tmp_dir}' EXIT
+
+for whl_file in "$@"; do
+    whl_file=$(realpath "${whl_file}")
+    whl_dir="${tmp_dir}/$(basename "${whl_file}")_unzipped"
+    mkdir -pv "${whl_dir}"
+    (
+        set -x
+        unzip -q "${whl_file}" -d "${whl_dir}"
+    )
+    version_with_suffix=$(grep '^Version:' "${whl_dir}"/*/METADATA | cut -d' ' -f2)
+    version_with_suffix_escaped=${version_with_suffix/+/%2B}
+    # Remove all suffixed +bleh versions
+    version_no_suffix=${version_with_suffix/+*/}
+    new_whl_file=${whl_file/${version_with_suffix_escaped}/${version_no_suffix}}
+    dist_info_folder=$(find "${whl_dir}" -type d -name '*.dist-info' | head -1)
+    basename_dist_info_folder=$(basename "${dist_info_folder}")
+    dirname_dist_info_folder=$(dirname "${dist_info_folder}")
+    (
+        set -x
+        find "${dist_info_folder}" -type f -exec sed -i "s!${version_with_suffix}!${version_no_suffix}!" {} \;
+        # Moves distinfo from one with a version suffix to one without
+        # Example: torch-1.8.0+cpu.dist-info => torch-1.8.0.dist-info
+        mv "${dist_info_folder}" "${dirname_dist_info_folder}/${basename_dist_info_folder/${version_with_suffix}/${version_no_suffix}}"
+        cd "${whl_dir}"
+        zip -qr "${new_whl_file}" .
+    )
+done