From 93c5a5d70f0f99b880552f6fd869eeec6afc7942 Mon Sep 17 00:00:00 2001
From: Michael Sarahan <msarahan@nvidia.com>
Date: Wed, 3 Apr 2024 10:04:34 -0700
Subject: [PATCH] do not require first arg to upload to be pkg_type

---
 tools/rapids-download-wheels-from-s3 | 15 +++++++--------
 tools/rapids-upload-wheels-to-s3     | 15 +++++++--------
 2 files changed, 14 insertions(+), 16 deletions(-)

diff --git a/tools/rapids-download-wheels-from-s3 b/tools/rapids-download-wheels-from-s3
index c58a6fe..be6ea1b 100755
--- a/tools/rapids-download-wheels-from-s3
+++ b/tools/rapids-download-wheels-from-s3
@@ -6,15 +6,14 @@ set -eo pipefail
 
 source rapids-prompt-local-repo-config
 
+# For legacy reasons, allow this script to be run without the pkg_type being the first arg
+pkg_name="$(rapids-package-name "wheel_python")"
+
 pkg_type="$1"
-if [ "${pkg_type}" != "cpp" ] && [ "${pkg_type}" != "python" ]; then
-    rapids-echo-error "Wheel cpp packages are now supported."
-    rapids-echo-error "You must specify package type as the first argument to rapids-upload-wheels-to-s3"
+if [ "${pkg_type}" = "cpp" ] && [ "${pkg_type}" = "python" ]; then
+    # remove pkg_type from args because we handle it in this script
+    shift;
+    pkg_name="$(rapids-package-name "wheel_${pkg_type}")"
 fi
 
-# remove pkg_type from args because we handle it in this script
-shift;
-
-pkg_name="$(rapids-package-name "wheel_${pkg_type}")"
-
 rapids-download-from-s3 "${pkg_name}" "$@"
diff --git a/tools/rapids-upload-wheels-to-s3 b/tools/rapids-upload-wheels-to-s3
index 9e35774..24e1548 100755
--- a/tools/rapids-upload-wheels-to-s3
+++ b/tools/rapids-upload-wheels-to-s3
@@ -4,21 +4,20 @@
 #   1) wheel path to tar up
 set -e
 
+# For legacy reasons, allow this script to be run without the pkg_type being the first arg
+pkg_name="$(rapids-package-name "wheel_python")"
+
 pkg_type="$1"
-if [ "${pkg_type}" != "cpp" ] && [ "${pkg_type}" != "python" ]; then
-    rapids-echo-error "Wheel cpp packages are now supported."
-    rapids-echo-error "You must specify package type as the first argument to rapids-upload-wheels-to-s3"
+if [ "${pkg_type}" = "cpp" ] && [ "${pkg_type}" = "python" ]; then
+    # remove pkg_type from args because we handle it in this script
+    shift;
+    pkg_name="$(rapids-package-name "wheel_${pkg_type}")"
 fi
 
-# remove pkg_type from args because we handle it in this script
-shift;
-
 if [ "${CI:-false}" = "false" ]; then
   rapids-echo-stderr "Packages from local builds cannot be uploaded to S3."
   rapids-echo-stderr "Open a PR to have successful builds uploaded."
   exit 0
 fi
 
-pkg_name="$(rapids-package-name "wheel_${pkg_type}")"
-
 rapids-upload-to-s3 "${pkg_name}" "$@"