feat #1342: support specify custom python3 venv

2026-09-05 23:51:25 +08:00 · 2022-06-13 21:31:05 +08:00
parent abd31a8d88
commit fac6dc4955
35 changed files with 381 additions and 123 deletions
@@ -13,12 +13,13 @@ import (
 	"strconv"
 	"strings"

-	"github.com/httprunner/funplugin/shared"
+	"github.com/httprunner/funplugin/fungo"
 	"github.com/pkg/errors"
 	"github.com/rs/zerolog/log"
 	"gopkg.in/yaml.v3"

 	"github.com/httprunner/httprunner/v4/hrp/internal/json"
+	"github.com/httprunner/httprunner/v4/hrp/internal/version"
 )

 func Dump2JSON(data interface{}, path string) error {
@@ -88,24 +89,104 @@ func FormatResponse(raw interface{}) interface{} {
 	return formattedResponse
 }

-func EnsurePython3Venv(packages ...string) (string, error) {
-	// create python venv
+var Python3Executable string = "python3" // system default python3
+
+func PrepareVenv(venv string) error {
+	defer func() {
+		log.Info().Str("Python3Executable", Python3Executable).Msg("set python3 executable path")
+	}()
+
+	// specify python3 venv
+	if venv != "" {
+		Python3Executable = getPython3Executable(venv)
+		return nil
+	}
+
+	// not specify python3 venv, create default venv
+	python3, err := createDefaultPython3Venv()
+	if err != nil {
+		return errors.Wrap(err, "create default python3 venv failed")
+	}
+	Python3Executable = python3
+
+	return nil
+}
+
+// create default python3 venv located in $HOME/.hrp/venv
+func createDefaultPython3Venv() (string, error) {
+	log.Info().Msg("create default python3 venv at $HOME/.hrp/venv")
 	home, err := os.UserHomeDir()
 	if err != nil {
 		return "", errors.Wrap(err, "get user home dir failed")
 	}
-	venvDir := filepath.Join(home, ".hrp", "venv")
-	python3, err := shared.EnsurePython3Venv(venvDir, packages...)
-	if err != nil {
-		return "", errors.Wrap(err, "ensure python3 venv failed")
+	venv := filepath.Join(home, ".hrp", "venv")
+	packages := []string{
+		fmt.Sprintf("funppy==%s", fungo.Version),
+		fmt.Sprintf("httprunner==%s", version.VERSION),
+	}
+	return EnsurePython3Venv(venv, packages...)
+}
+
+func ExecPython3Command(cmdName string, args ...string) error {
+	args = append([]string{"-m", cmdName}, args...)
+	return ExecCommand(Python3Executable, args...)
+}
+
+func InstallPythonPackage(python3 string, pkg string) (err error) {
+	var pkgName string
+	if strings.Contains(pkg, "==") {
+		// funppy==0.4.2
+		pkgInfo := strings.Split(pkg, "==")
+		pkgName = pkgInfo[0]
+	} else if strings.Contains(pkg, ">=") {
+		// httprunner>=4.0.0-beta
+		pkgInfo := strings.Split(pkg, ">=")
+		pkgName = pkgInfo[0]
+	} else {
+		pkgName = pkg
 	}

-	return python3, nil
+	defer func() {
+		if err == nil {
+			// check package version
+			if out, err := exec.Command(
+				python3, "-c", fmt.Sprintf("import %s; print(%s.__version__)", pkgName, pkgName),
+			).Output(); err == nil {
+				log.Info().
+					Str("name", pkgName).
+					Str("version", strings.TrimSpace(string(out))).
+					Msg("python package is ready")
+			}
+		}
+	}()
+
+	// check if package installed
+	err = exec.Command(python3, "-c", fmt.Sprintf("import %s", pkgName)).Run()
+	if err == nil {
+		return nil
+	}
+
+	// check if pip available
+	err = ExecCommand(python3, "-m", "pip", "--version")
+	if err != nil {
+		log.Warn().Msg("pip is not available")
+		return errors.Wrap(err, "pip is not available")
+	}
+
+	log.Info().Str("package", pkg).Msg("installing python package")
+
+	// install package
+	err = ExecCommand(python3, "-m", "pip", "install", "--upgrade", pkg,
+		"--quiet", "--disable-pip-version-check")
+	if err != nil {
+		return errors.Wrap(err, "pip install package failed")
+	}
+
+	return nil
 }

 func CheckPythonScriptSyntax(path string) error {
-	err := ExecCommand("python3", "-m", "py_compile", path)
-	return err
+	return ExecCommand("python3", "-m", "py_compile", path)
 }

 func ExecCommandInDir(cmd *exec.Cmd, dir string) error {
@@ -125,30 +206,6 @@ func ExecCommandInDir(cmd *exec.Cmd, dir string) error {
 	return nil
 }

-func ExecCommand(cmdName string, args ...string) error {
-	cmd := exec.Command(cmdName, args...)
-	log.Info().Str("cmd", cmd.String()).Msg("exec command")
-
-	// print output with colors
-	cmd.Stdout = os.Stdout
-	cmd.Stderr = os.Stderr
-
-	// add cmd dir path to PATH
-	PATH := fmt.Sprintf("%s:%s", filepath.Dir(cmdName), os.Getenv("PATH"))
-	if err := os.Setenv("PATH", PATH); err != nil {
-		log.Error().Err(err).Msg("failed to add cmd dir path to $PATH")
-		return err
-	}
-
-	err := cmd.Run()
-	if err != nil {
-		log.Error().Err(err).Msg("exec command failed")
-		return err
-	}
-
-	return err
-}
-
 func CreateFolder(folderPath string) error {
 	log.Info().Str("path", folderPath).Msg("create folder")
 	err := os.MkdirAll(folderPath, os.ModePerm)
@@ -0,0 +1,77 @@
+// +build darwin linux
+
+package builtin
+
+import (
+	"fmt"
+	"os"
+	"os/exec"
+	"path/filepath"
+
+	"github.com/pkg/errors"
+	"github.com/rs/zerolog/log"
+)
+
+func getPython3Executable(venvDir string) string {
+	return filepath.Join(venvDir, "bin", "python3")
+}
+
+// EnsurePython3Venv ensures python3 venv for hashicorp python plugin
+// venvDir should be directory path of target venv
+func EnsurePython3Venv(venvDir string, packages ...string) (python3 string, err error) {
+	python3 = getPython3Executable(venvDir)
+
+	log.Info().
+		Str("python3", python3).
+		Interface("packages", packages).
+		Msg("ensure python3 venv")
+
+	// check if python3 venv is available
+	if err := exec.Command(python3, "--version").Run(); err != nil {
+		// python3 venv not available, create one
+		// check if system python3 is available
+		if err := ExecCommand("python3", "--version"); err != nil {
+			return "", errors.Wrap(err, "python3 not found")
+		}
+
+		// check if .venv exists
+		if _, err := os.Stat(venvDir); err == nil {
+			// .venv exists, remove first
+			if err := ExecCommand("rm", "-rf", venvDir); err != nil {
+				return "", errors.Wrap(err, "remove existed venv failed")
+			}
+		}
+
+		// create python3 .venv
+		if err := ExecCommand("python3", "-m", "venv", venvDir); err != nil {
+			return "", errors.Wrap(err, "create python3 venv failed")
+		}
+	}
+
+	// install default python packages
+	for _, pkg := range packages {
+		err := InstallPythonPackage(python3, pkg)
+		if err != nil {
+			return "", errors.Wrap(err, fmt.Sprintf("pip install %s failed", pkg))
+		}
+	}
+
+	return python3, nil
+}
+
+func ExecCommand(cmdName string, args ...string) error {
+	cmd := exec.Command(cmdName, args...)
+	log.Info().Str("cmd", cmd.String()).Msg("exec command")
+
+	// print output with colors
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+
+	err := cmd.Run()
+	if err != nil {
+		log.Error().Err(err).Msg("exec command failed")
+		return err
+	}
+
+	return nil
+}
@@ -0,0 +1,100 @@
+// +build windows
+
+package builtin
+
+import (
+	"fmt"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strings"
+
+	"github.com/pkg/errors"
+	"github.com/rs/zerolog/log"
+)
+
+func getPython3Executable(venvDir string) string {
+	return filepath.Join(venvDir, "Scripts", "python3.exe")
+}
+
+// EnsurePython3Venv ensures python3 venv for hashicorp python plugin
+// venvDir should be directory path of target venv
+func EnsurePython3Venv(venvDir string, packages ...string) (python3 string, err error) {
+	python3 = getPython3Executable(venvDir)
+	log.Info().
+		Str("python3", python3).
+		Interface("packages", packages).
+		Msg("ensure python3 venv")
+
+	systemPython := "python3"
+
+	// check if python3 venv is available
+	if err := exec.Command("cmd", "/c", python3, "--version").Run(); err != nil {
+		// python3 venv not available, create one
+		// check if system python3 is available
+		if err := ExecCommand(systemPython, "--version"); err != nil {
+			if err := ExecCommand("python", "--version"); err != nil {
+				return "", errors.Wrap(err, "python3 not found")
+			}
+			systemPython = "python"
+		}
+
+		// check if .venv exists
+		if _, err := os.Stat(venvDir); err == nil {
+			// .venv exists, remove first
+			if err := ExecCommand("del", "/q", venvDir); err != nil {
+				return "", errors.Wrap(err, "remove existed venv failed")
+			}
+		}
+
+		// create python3 .venv
+		// notice: --symlinks should be specified for windows
+		// https://github.com/actions/virtual-environments/issues/2690
+		if err := ExecCommand(systemPython, "-m", "venv", "--symlinks", venvDir); err != nil {
+			// fix: failed to symlink on Windows
+			log.Warn().Msg("failed to create python3 .venv by using --symlinks, try to use --copies")
+			if err := ExecCommand(systemPython, "-m", "venv", "--copies", venvDir); err != nil {
+				return "", errors.Wrap(err, "create python3 venv failed")
+			}
+		}
+
+		// fix: python3 doesn't exist in .venv on Windows
+		if _, err := os.Stat(python3); err != nil {
+			log.Warn().Msg("python3 doesn't exist, try to link python")
+			err := os.Link(filepath.Join(venvDir, "Scripts", "python.exe"), python3)
+			if err != nil {
+				return "", errors.Wrap(err, "python3 doesn't exist in .venv")
+			}
+		}
+	}
+
+	// install default python packages
+	for _, pkg := range packages {
+		err := InstallPythonPackage(python3, pkg)
+		if err != nil {
+			return "", errors.Wrap(err, fmt.Sprintf("pip install %s failed", pkg))
+		}
+	}
+
+	return python3, nil
+}
+
+func ExecCommand(cmdName string, args ...string) error {
+	// "cmd /c" carries out the command specified by string and then stops
+	// refer: https://docs.microsoft.com/en-us/windows-server/administration/windows-commands/cmd
+	cmdStr := fmt.Sprintf("%s %s", cmdName, strings.Join(args, " "))
+	cmd := exec.Command("cmd", "/c", cmdStr)
+	log.Info().Str("cmd", cmd.String()).Msg("exec command")
+
+	// print output with colors
+	cmd.Stdout = os.Stdout
+	cmd.Stderr = os.Stderr
+
+	err := cmd.Run()
+	if err != nil {
+		log.Error().Err(err).Msg("exec command failed")
+		return err
+	}
+
+	return nil
+}