From 4cbbb57e1f0ff425f66350621d1c20390825d9e4 Mon Sep 17 00:00:00 2001
From: Amin Yahyaabadi <aminyahyaabadi74@gmail.com>
Date: Tue, 4 Jul 2023 12:29:30 -0700
Subject: [PATCH] feat: upgrade setup-tools and refactor package installer

---
 src/python/python.ts            | 25 +++++++++++--------------
 src/utils/setup/setupPipPack.ts | 33 ++++++++++++++++++++++++---------
 2 files changed, 35 insertions(+), 23 deletions(-)

diff --git a/src/python/python.ts b/src/python/python.ts
index b2c861fe..6f441dd4 100644
--- a/src/python/python.ts
+++ b/src/python/python.ts
@@ -3,7 +3,7 @@ import { getExecOutput } from "@actions/exec"
 import assert from "assert"
 import { GITHUB_ACTIONS } from "ci-info"
 import { info, warning } from "ci-log"
-import { execaSync } from "execa"
+import { execa } from "execa"
 import memoize from "micro-memoize"
 import { addExeExt, dirname, join } from "patha"
 import which from "which"
@@ -21,6 +21,7 @@ import { isBinUptoDate } from "../utils/setup/version"
 import { unique } from "../utils/std"
 import { MinVersions } from "../versions/default_versions"
 import { pathExists } from "path-exists"
+import { setupPipPackWithPython } from "../utils/setup/setupPipPack"
 
 export async function setupPython(version: string, setupDir: string, arch: string): Promise<InstallationInfo> {
   const installInfo = await findOrSetupPython(version, setupDir, arch)
@@ -33,11 +34,12 @@ export async function setupPython(version: string, setupDir: string, arch: strin
     throw new Error("pip was not installed correctly")
   }
 
-  // setup wheel
+  // setup wheel and setuptools
   try {
-    setupWheel(foundPython)
+    await setupPipPackWithPython(foundPython, "setuptools", undefined, true)
+    await setupPipPackWithPython(foundPython, "wheel", undefined, true)
   } catch (err) {
-    warning(`Failed to install wheels: ${(err as Error).toString()}. Ignoring...`)
+    warning(`Failed to install setuptools or wheel: ${(err as Error).toString()}. Ignoring...`)
   }
 
   return installInfo
@@ -194,23 +196,23 @@ async function isPipUptoDate(pip: string) {
 }
 
 async function setupPip(foundPython: string) {
-  const upgraded = ensurePipUpgrade(foundPython)
+  const upgraded = await ensurePipUpgrade(foundPython)
   if (!upgraded) {
     await setupPipSystem()
     // upgrade pip
-    ensurePipUpgrade(foundPython)
+    await ensurePipUpgrade(foundPython)
   }
 }
 
-function ensurePipUpgrade(foundPython: string) {
+async function ensurePipUpgrade(foundPython: string) {
   try {
-    execaSync(foundPython, ["-m", "ensurepip", "-U", "--upgrade"], { stdio: "inherit" })
+    await execa(foundPython, ["-m", "ensurepip", "-U", "--upgrade"], { stdio: "inherit" })
     return true
   } catch (err1) {
     info((err1 as Error)?.toString?.())
     try {
       // ensure pip is disabled on Ubuntu
-      execaSync(foundPython, ["-m", "pip", "install", "--upgrade", "pip"], { stdio: "inherit" })
+      await execa(foundPython, ["-m", "pip", "install", "--upgrade", "pip"], { stdio: "inherit" })
       return true
     } catch (err2) {
       info((err2 as Error)?.toString?.())
@@ -235,11 +237,6 @@ function setupPipSystem() {
   throw new Error(`Could not install pip on ${process.platform}`)
 }
 
-/** Install wheel (required for Conan, Meson, etc.) */
-function setupWheel(foundPython: string) {
-  execaSync(foundPython, ["-m", "pip", "install", "-U", "wheel"], { stdio: "inherit" })
-}
-
 async function addPythonBaseExecPrefix_raw(python: string) {
   const dirs: string[] = []
 
diff --git a/src/utils/setup/setupPipPack.ts b/src/utils/setup/setupPipPack.ts
index 87d33dc6..517387ff 100644
--- a/src/utils/setup/setupPipPack.ts
+++ b/src/utils/setup/setupPipPack.ts
@@ -8,23 +8,29 @@ import { addPath } from "../env/addEnv"
 import { InstallationInfo } from "./setupBin"
 import { getVersion } from "../../versions/versions"
 import { ubuntuVersion } from "../env/ubuntu_version"
-
-/* eslint-disable require-atomic-updates */
-let python: string | undefined
+import memoize from "micro-memoize"
 
 /** A function that installs a package using pip */
-export async function setupPipPack(name: string, version?: string): Promise<InstallationInfo> {
+export async function setupPipPack(name: string, version?: string, upgrade = false): Promise<InstallationInfo> {
+  return setupPipPackWithPython(await getPython(), name, version, upgrade)
+}
+
+export async function setupPipPackWithPython(
+  givenPython: string,
+  name: string,
+  version?: string,
+  upgrade = false
+): Promise<InstallationInfo> {
   info(`Installing ${name} ${version ?? ""} via pip`)
 
-  if (python === undefined) {
-    python = (await setupPython(getVersion("python", undefined, await ubuntuVersion()), "", process.arch)).bin!
-  }
+  const nameAndVersion = version !== undefined && version !== "" ? `${name}==${version}` : name
+  const upgradeFlag = upgrade === true ? ["--upgrade"] : []
 
-  execaSync(python, ["-m", "pip", "install", version !== undefined && version !== "" ? `${name}==${version}` : name], {
+  execaSync(givenPython, ["-m", "pip", "install", ...upgradeFlag, nameAndVersion], {
     stdio: "inherit",
   })
 
-  const execPaths = await addPythonBaseExecPrefix(python)
+  const execPaths = await addPythonBaseExecPrefix(givenPython)
   const binDir = await findBinDir(execPaths, name)
 
   await addPath(binDir)
@@ -32,6 +38,15 @@ export async function setupPipPack(name: string, version?: string): Promise<Inst
   return { binDir }
 }
 
+async function getPython_raw(): Promise<string> {
+  const pythonBin = (await setupPython(getVersion("python", undefined, await ubuntuVersion()), "", process.arch)).bin
+  if (pythonBin === undefined) {
+    throw new Error("Python binary was not found")
+  }
+  return pythonBin
+}
+const getPython = memoize(getPython_raw)
+
 async function findBinDir(dirs: string[], name: string) {
   const exists = await Promise.all(dirs.map((dir) => pathExists(join(dir, addExeExt(name)))))
   const dirIndex = exists.findIndex((exist) => exist)