Update main.py

Author: Jason2866

builder/main.py

@@ -36,6 +36,29 @@
 from platformio.util import get_serial_ports
 
 
+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)
+        if normalized_path not in current_paths:
+            os.environ["PYTHONPATH"] += os.pathsep + normalized_path
+    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)
+
+
 # Initialize environment and configuration
 env = DefaultEnvironment()
 platform = env.PioPlatform()
@@ -45,6 +68,13 @@
 # Framework directory path
 FRAMEWORK_DIR = platform.get_package_dir("framework-arduinoespressif32")
 
+# Add framework Python tools to path if available
+if FRAMEWORK_DIR:
+    framework_python_path = os.path.join(FRAMEWORK_DIR, "tools", "python")
+    if os.path.isdir(framework_python_path):
+        add_to_pythonpath(framework_python_path)
+
+# Python dependencies required for the build process
 python_deps = {
     "uv": ">=0.1.0",
     "pyyaml": ">=6.0.2",
@@ -57,7 +87,16 @@
 
 
 def get_packages_to_install(deps, installed_packages):
-    """Generator for Python packages to install"""
+    """
+    Generator for Python packages that need to be installed.
+    
+    Args:
+        deps (dict): Dictionary of package names and version specifications
+        installed_packages (dict): Dictionary of currently installed packages
+        
+    Yields:
+        str: Package name that needs to be installed
+    """
     for package, spec in deps.items():
         if package not in installed_packages:
             yield package
@@ -68,7 +107,12 @@ def get_packages_to_install(deps, installed_packages):
 
 
 def install_python_deps():
-    """Ensure uv package manager is available, install with pip if not"""
+    """
+    Ensure uv package manager is available and install required Python dependencies.
+    
+    Returns:
+        bool: True if successful, False otherwise
+    """
     try:
         result = subprocess.run(
             ["uv", "--version"],
@@ -86,7 +130,8 @@ def install_python_deps():
                 [env.subst("$PYTHONEXE"), "-m", "pip", "install", "uv>=0.1.0", "-q", "-q", "-q"],
                 capture_output=True,
                 text=True,
-                timeout=30  # 30 second timeout
+                timeout=30,  # 30 second timeout
+                env=os.environ  # Use modified environment with custom PYTHONPATH
             )
             if result.returncode != 0:
                 if result.stderr:
@@ -104,6 +149,12 @@ def install_python_deps():
 
 
     def _get_installed_uv_packages():
+        """
+        Get list of installed packages using uv.
+        
+        Returns:
+            dict: Dictionary of installed packages with versions
+        """
         result = {}
         try:
             cmd = ["uv", "pip", "list", "--format=json"]
@@ -112,7 +163,8 @@ def _get_installed_uv_packages():
                 capture_output=True,
                 text=True,
                 encoding='utf-8',
-                timeout=30  # 30 second timeout
+                timeout=30,  # 30 second timeout
+                env=os.environ  # Use modified environment with custom PYTHONPATH
             )
 
             if result_obj.returncode == 0:
@@ -154,7 +206,8 @@ def _get_installed_uv_packages():
                 cmd,
                 capture_output=True,
                 text=True,
-                timeout=30  # 30 second timeout for package installation
+                timeout=30,  # 30 second timeout for package installation
+                env=os.environ  # Use modified environment with custom PYTHONPATH
             )
 
             if result.returncode != 0:
@@ -177,10 +230,22 @@ def _get_installed_uv_packages():
 
 
 def install_esptool(env):
-    """Install esptool from package folder "tool-esptoolpy" using uv package manager"""
+    """
+    Install esptool from package folder "tool-esptoolpy" using uv package manager.
+    
+    Args:
+        env: SCons environment object
+        
+    Returns:
+        bool: True if successful, False otherwise
+    """
     try:
-        subprocess.check_call([env.subst("$PYTHONEXE"), "-c", "import esptool"], 
-                              stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
+        subprocess.check_call(
+            [env.subst("$PYTHONEXE"), "-c", "import esptool"], 
+            stdout=subprocess.DEVNULL, 
+            stderr=subprocess.DEVNULL,
+            env=os.environ  # Use modified environment with custom PYTHONPATH
+        )
         return True
     except (subprocess.CalledProcessError, FileNotFoundError):
         pass
@@ -192,7 +257,7 @@ def install_esptool(env):
                 "uv", "pip", "install", "--quiet",
                 f"--python={env.subst('$PYTHONEXE')}",
                 "-e", esptool_repo_path
-            ])
+            ], env=os.environ)  # Use modified environment with custom PYTHONPATH
             return True
         except subprocess.CalledProcessError as e:
             print(f"Warning: Failed to install esptool: {e}")
@@ -201,6 +266,7 @@ def install_esptool(env):
     return False
 
 
+# Install Python dependencies and esptool
 install_python_deps()
 install_esptool(env)
 
@@ -209,6 +275,11 @@ def BeforeUpload(target, source, env):
     """
     Prepare the environment before uploading firmware.
     Handles port detection and special upload configurations.
+    
+    Args:
+        target: SCons target
+        source: SCons source
+        env: SCons environment object
     """
     upload_options = {}
     if "BOARD" in env:
@@ -228,7 +299,12 @@ def BeforeUpload(target, source, env):
 def _get_board_memory_type(env):
     """
     Determine the memory type configuration for the board.
-    Returns the appropriate memory type string based on board configuration.
+    
+    Args:
+        env: SCons environment object
+        
+    Returns:
+        str: The appropriate memory type string based on board configuration
     """
     board_config = env.BoardConfig()
     default_type = "%s_%s" % (
@@ -249,20 +325,41 @@ def _get_board_memory_type(env):
 def _normalize_frequency(frequency):
     """
     Convert frequency value to normalized string format (e.g., "40m").
-    Removes 'L' suffix and converts to MHz format.
+    
+    Args:
+        frequency: Frequency value to normalize
+        
+    Returns:
+        str: Normalized frequency string with 'm' suffix
     """
     frequency = str(frequency).replace("L", "")
     return str(int(int(frequency) / 1000000)) + "m"
 
 
 def _get_board_f_flash(env):
-    """Get the flash frequency for the board."""
+    """
+    Get the flash frequency for the board.
+    
+    Args:
+        env: SCons environment object
+        
+    Returns:
+        str: Flash frequency string
+    """
     frequency = env.subst("$BOARD_F_FLASH")
     return _normalize_frequency(frequency)
 
 
 def _get_board_f_image(env):
-    """Get the image frequency for the board, fallback to flash frequency."""
+    """
+    Get the image frequency for the board, fallback to flash frequency.
+    
+    Args:
+        env: SCons environment object
+        
+    Returns:
+        str: Image frequency string
+    """
     board_config = env.BoardConfig()
     if "build.f_image" in board_config:
         return _normalize_frequency(board_config.get("build.f_image"))
@@ -271,7 +368,15 @@ def _get_board_f_image(env):
 
 
 def _get_board_f_boot(env):
-    """Get the boot frequency for the board, fallback to flash frequency."""
+    """
+    Get the boot frequency for the board, fallback to flash frequency.
+    
+    Args:
+        env: SCons environment object
+        
+    Returns:
+        str: Boot frequency string
+    """
     board_config = env.BoardConfig()
     if "build.f_boot" in board_config:
         return _normalize_frequency(board_config.get("build.f_boot"))
@@ -283,6 +388,12 @@ def _get_board_flash_mode(env):
     """
     Determine the appropriate flash mode for the board.
     Handles special cases for OPI memory types.
+    
+    Args:
+        env: SCons environment object
+        
+    Returns:
+        str: Flash mode string
     """
     if _get_board_memory_type(env) in ("opi_opi", "opi_qspi"):
         return "dout"
@@ -297,6 +408,12 @@ def _get_board_boot_mode(env):
     """
     Determine the boot mode for the board.
     Handles special cases for OPI memory types.
+    
+    Args:
+        env: SCons environment object
+        
+    Returns:
+        str: Boot mode string
     """
     memory_type = env.BoardConfig().get("build.arduino.memory_type", "")
     build_boot = env.BoardConfig().get("build.boot", "$BOARD_FLASH_MODE")
@@ -308,7 +425,12 @@ def _get_board_boot_mode(env):
 def _parse_size(value):
     """
     Parse size values from various formats (int, hex, K/M suffixes).
-    Returns the size in bytes as an integer.
+    
+    Args:
+        value: Size value to parse
+        
+    Returns:
+        int: Size in bytes as an integer
     """
     if isinstance(value, int):
         return value
@@ -326,6 +448,12 @@ def _parse_partitions(env):
     """
     Parse the partition table CSV file and return partition information.
     Also sets the application offset for the environment.
+    
+    Args:
+        env: SCons environment object
+        
+    Returns:
+        list: List of partition dictionaries
     """
     partitions_csv = env.subst("$PARTITIONS_TABLE_CSV")
     if not isfile(partitions_csv):
@@ -377,6 +505,9 @@ def _update_max_upload_size(env):
     """
     Update the maximum upload size based on partition table configuration.
     Prioritizes user-specified partition names.
+    
+    Args:
+        env: SCons environment object
     """
     if not env.get("PARTITIONS_TABLE_CSV"):
         return
@@ -412,14 +543,25 @@ def _update_max_upload_size(env):
 
 
 def _to_unix_slashes(path):
-    """Convert Windows-style backslashes to Unix-style forward slashes."""
+    """
+    Convert Windows-style backslashes to Unix-style forward slashes.
+    
+    Args:
+        path (str): Path to convert
+        
+    Returns:
+        str: Path with Unix-style slashes
+    """
     return path.replace("\\", "/")
 
 
 def fetch_fs_size(env):
     """
     Extract filesystem size and offset information from partition table.
     Sets FS_START, FS_SIZE, FS_PAGE, and FS_BLOCK environment variables.
+    
+    Args:
+        env: SCons environment object
     """
     fs = None
     for p in _parse_partitions(env):
@@ -450,15 +592,27 @@ def fetch_fs_size(env):
 
 
 def __fetch_fs_size(target, source, env):
-    """Wrapper function for fetch_fs_size to be used as SCons emitter."""
+    """
+    Wrapper function for fetch_fs_size to be used as SCons emitter.
+    
+    Args:
+        target: SCons target
+        source: SCons source
+        env: SCons environment object
+        
+    Returns:
+        tuple: (target, source) tuple
+    """
     fetch_fs_size(env)
     return (target, source)
 
 
 def check_lib_archive_exists():
     """
     Check if lib_archive is set in platformio.ini configuration.
-    Returns True if found, False otherwise.
+    
+    Returns:
+        bool: True if found, False otherwise
     """
     for section in projectconfig.sections():
         if "lib_archive" in projectconfig.options(section):
@@ -611,8 +765,13 @@ def check_lib_archive_exists():
 
 def firmware_metrics(target, source, env):
     """
-    Custom target to run esp-idf-size with support for command line parameters
+    Custom target to run esp-idf-size with support for command line parameters.
     Usage: pio run -t metrics -- [esp-idf-size arguments]
+    
+    Args:
+        target: SCons target
+        source: SCons source
+        env: SCons environment object
     """
     if terminal_cp != "utf-8":
         print("Firmware metrics can not be shown. Set the terminal codepage to \"utf-8\"")
@@ -655,8 +814,8 @@ def firmware_metrics(target, source, env):
         if env.GetProjectOption("custom_esp_idf_size_verbose", False):
             print(f"Running command: {' '.join(cmd)}")
 
-        # Call esp-idf-size
-        result = subprocess.run(cmd, check=False, capture_output=False)
+        # Call esp-idf-size with modified environment
+        result = subprocess.run(cmd, check=False, capture_output=False, env=os.environ)
 
         if result.returncode != 0:
             print(f"Warning: esp-idf-size exited with code {result.returncode}")