simplify venv setup (#246)

Author: Jason2866

builder/main.py

@@ -16,6 +16,7 @@
 import json
 import os
 import re
+import site
 import semantic_version
 import shlex
 import subprocess
@@ -61,19 +62,23 @@
 platform = env.PioPlatform()
 projectconfig = env.GetProjectConfig()
 terminal_cp = locale.getpreferredencoding().lower()
-PYTHON_EXE = env.subst("$PYTHONEXE")  # Global Python executable path
-
-# Framework directory path
 FRAMEWORK_DIR = platform.get_package_dir("framework-arduinoespressif32")
-
 platformio_dir = projectconfig.get("platformio", "core_dir")
+
+# Global Python executable path, replaced later with venv python path
+PYTHON_EXE = env.subst("$PYTHONEXE")
 penv_dir = os.path.join(platformio_dir, "penv")
 
-pip_path = os.path.join(
-    penv_dir,
-    "Scripts" if IS_WINDOWS else "bin",
-    "pip" + (".exe" if IS_WINDOWS else ""),
-)
+
+def get_executable_path(executable_name):
+    """
+    Get the path to an executable based on the penv_dir.
+    """
+    exe_suffix = ".exe" if IS_WINDOWS else ""
+    scripts_dir = "Scripts" if IS_WINDOWS else "bin"
+    
+    return os.path.join(penv_dir, scripts_dir, f"{executable_name}{exe_suffix}")
+
 
 def setup_pipenv_in_package():
     """
@@ -87,144 +92,43 @@ def setup_pipenv_in_package():
             )
         )
         assert os.path.isfile(
-            pip_path
+            get_executable_path("pip")
         ), "Error: Failed to create a proper virtual environment. Missing the `pip` binary!"
 
-    penv_python = os.path.join(penv_dir, "Scripts", "python.exe") if IS_WINDOWS else os.path.join(penv_dir, "bin", "python")
-    env.Replace(PYTHONEXE=penv_python)
 
-# Setup virtual environment if needed and find path to Python exe
+# Setup virtual environment if needed
 setup_pipenv_in_package()
+
 # Set Python Scons Var to env Python
-PYTHON_EXE = env.subst("$PYTHONEXE")
-# Remove PYTHONHOME if set
-os.environ.pop('PYTHONHOME', None)
+penv_python = get_executable_path("python")
+env.Replace(PYTHONEXE=penv_python)
+PYTHON_EXE = penv_python
 
 # check for python binary, exit with error when not found
 assert os.path.isfile(PYTHON_EXE), f"Python executable not found: {PYTHON_EXE}"
 
-def add_to_pythonpath(path):
-    """
-    Add a path to the PYTHONPATH environment variable (cross-platform).
-    
-    Args:
-        path (str): The path to add to PYTHONPATH
-    """
-    # Normalize the path for the current OS
-    normalized_path = os.path.normpath(path)
-
-    # Add to PYTHONPATH environment variable
-    if "PYTHONPATH" in os.environ:
-        current_paths = os.environ["PYTHONPATH"].split(os.pathsep)
-        normalized_current_paths = [os.path.normpath(p) for p in current_paths]
-        if normalized_path not in normalized_current_paths:
-            # Rebuild PYTHONPATH with normalized paths to avoid duplicates
-            normalized_current_paths.insert(0, normalized_path)
-            os.environ["PYTHONPATH"] = os.pathsep.join(normalized_current_paths)
-    else:
-        os.environ["PYTHONPATH"] = normalized_path
-
-    # Also add to sys.path for immediate availability
-    if normalized_path not in sys.path:
-        sys.path.insert(0, normalized_path)
-
 
 def setup_python_paths():
-    """
-    Setup Python paths based on the actual Python executable being used.
-    
-    This function configures both PYTHONPATH environment variable and sys.path
-    to include the Python executable directory and site-packages directory.
-    """    
-    # Get the directory containing the Python executable
-    python_dir = os.path.dirname(PYTHON_EXE)
-
-    # Add Scripts directory to PATH for Windows
-    if IS_WINDOWS:
-        scripts_dir = os.path.join(python_dir, "Scripts")
-        if os.path.isdir(scripts_dir):
-            os.environ["PATH"] = scripts_dir + os.pathsep + os.environ.get("PATH", "")
-    else:
-        bin_dir = os.path.join(python_dir, "bin")
-        if os.path.isdir(bin_dir):
-            os.environ["PATH"] = bin_dir + os.pathsep + os.environ.get("PATH", "")
-
-    penv_site_packages = None
-    if python_dir not in sys.path:
-        add_to_pythonpath(python_dir)
-    if IS_WINDOWS:
-        penv_site_packages = os.path.join(penv_dir, "Lib", "site-packages")
-    else:
-        # Find the actual site-packages directory in the venv
-        penv_lib_dir = os.path.join(penv_dir, "lib")
-        if os.path.isdir(penv_lib_dir):
-            for python_version_dir in os.listdir(penv_lib_dir):
-                if python_version_dir.startswith("python"):
-                    penv_site_packages = os.path.join(penv_lib_dir, python_version_dir, "site-packages")
-                    break
-
-    if penv_site_packages and os.path.isdir(penv_site_packages) and penv_site_packages not in sys.path:
-        add_to_pythonpath(penv_site_packages)
-
-setup_python_paths()
-
-
-def _get_executable_path(python_exe, executable_name):
-    """
-    Get the path to an executable binary (esptool, uv, etc.) based on the Python executable path.
-    
-    Args:
-        python_exe (str): Path to Python executable
-        executable_name (str): Name of the executable to find (e.g., 'esptool', 'uv')
-        
-    Returns:
-        str: Path to executable or fallback to executable name
-    """
-    
-    python_dir = os.path.dirname(python_exe)
-    
-    if IS_WINDOWS:
-        executable_path = os.path.join(python_dir, f"{executable_name}.exe")
-    else:
-        # For Unix-like systems, executables are typically in the same directory as python
-        # or in a bin subdirectory
-        executable_path = os.path.join(python_dir, executable_name)
-        
-        # If not found in python directory, try bin subdirectory
-        if not os.path.isfile(executable_path):
-            bin_dir = os.path.join(python_dir, "bin")
-            executable_path = os.path.join(bin_dir, executable_name)
+    """Setup Python module search paths using the penv_dir."""    
+    # Add penv_dir to module search path
+    site.addsitedir(penv_dir)
 
-    if os.path.isfile(executable_path):
-        return executable_path
+    # Add site-packages directory
+    python_ver = f"python{sys.version_info.major}.{sys.version_info.minor}"
+    site_packages = (
+        os.path.join(penv_dir, "Lib", "site-packages") if IS_WINDOWS
+        else os.path.join(penv_dir, "lib", python_ver, "site-packages")
+    )
 
-    return executable_name  # Fallback to command name
+    if os.path.isdir(site_packages):
+        site.addsitedir(site_packages)
 
 
-def _get_esptool_executable_path(python_exe):
-    """
-    Get the path to the esptool executable binary.
-    
-    Args:
-        python_exe (str): Path to Python executable
-        
-    Returns:
-        str: Path to esptool executable
-    """
-    return _get_executable_path(python_exe, "esptool")
+setup_python_paths()
 
-
-def _get_uv_executable_path(python_exe):
-    """
-    Get the path to the uv executable binary.
-    
-    Args:
-        python_exe (str): Path to Python executable
-        
-    Returns:
-        str: Path to uv executable
-    """
-    return _get_executable_path(python_exe, "uv")
+# Set executable paths from tools
+esptool_binary_path = get_executable_path("esptool")
+uv_executable = get_executable_path("uv")
 
 
 def get_packages_to_install(deps, installed_packages):
@@ -254,9 +158,6 @@ def install_python_deps():
     Returns:
         bool: True if successful, False otherwise
     """
-    # Get uv executable path
-    uv_executable = _get_uv_executable_path(PYTHON_EXE)
-    
     try:
         result = subprocess.run(
             [uv_executable, "--version"],
@@ -275,15 +176,12 @@ def install_python_deps():
                 capture_output=True,
                 text=True,
                 timeout=30,  # 30 second timeout
-                env=os.environ  # Use modified environment with custom PYTHONPATH
+                env=os.environ  # Use current environment with venv Python
             )
             if result.returncode != 0:
                 if result.stderr:
                     print(f"Error output: {result.stderr.strip()}")
                 return False
-            
-            # Update uv executable path after installation
-            uv_executable = _get_uv_executable_path(PYTHON_EXE)
 
         except subprocess.TimeoutExpired:
             print("Error: uv installation timed out")
@@ -312,7 +210,7 @@ def _get_installed_uv_packages():
                 text=True,
                 encoding='utf-8',
                 timeout=30,  # 30 second timeout
-                env=os.environ  # Use modified environment with custom PYTHONPATH
+                env=os.environ  # Use current environment with venv Python
             )
 
             if result_obj.returncode == 0:
@@ -355,7 +253,7 @@ def _get_installed_uv_packages():
                 capture_output=True,
                 text=True,
                 timeout=30,  # 30 second timeout for package installation
-                env=os.environ  # Use modified environment with custom PYTHONPATH
+                env=os.environ  # Use current environment with venv Python
             )
 
             if result.returncode != 0:
@@ -380,9 +278,6 @@ def _get_installed_uv_packages():
 def install_esptool():
     """
     Install esptool from package folder "tool-esptoolpy" using uv package manager.
-    
-    Returns:
-        str: Path to esptool executable
 
     Raises:
         SystemExit: If esptool installation fails
@@ -394,35 +289,33 @@ def install_esptool():
             stderr=subprocess.DEVNULL,
             env=os.environ
         )
-        return _get_esptool_executable_path(PYTHON_EXE)
+        return
     except (subprocess.CalledProcessError, FileNotFoundError):
         pass
 
     esptool_repo_path = env.subst(platform.get_package_dir("tool-esptoolpy") or "")
     if not esptool_repo_path or not os.path.isdir(esptool_repo_path):
         print("Error: esptool package directory not found")
         sys.exit(1)
-        
-    uv_executable = _get_uv_executable_path(PYTHON_EXE)
+
     try:
         subprocess.check_call([
             uv_executable, "pip", "install", "--quiet",
             f"--python={PYTHON_EXE}",
             "-e", esptool_repo_path
         ], env=os.environ)
-        
-        return _get_esptool_executable_path(PYTHON_EXE)
-        
+
+        return
+
     except subprocess.CalledProcessError as e:
         print(f"Error: Failed to install esptool: {e}")
         sys.exit(1)
 
 
-# Install Python dependencies
+# Install espressif32 Python dependencies
 install_python_deps()
-
 # Install esptool after dependencies
-esptool_binary_path = install_esptool()
+install_esptool()
 
 
 def BeforeUpload(target, source, env):