|
|
@@ -1,4 +1,5 @@
|
|
|
import os
|
|
|
+import platform
|
|
|
import subprocess
|
|
|
from pathlib import Path
|
|
|
from typing import Dict, List, Optional
|
|
|
@@ -6,17 +7,32 @@ from typing import Dict, List, Optional
|
|
|
|
|
|
def extract_unix_vars_of_source(script_paths: List[Path]) -> Dict[str, str]:
|
|
|
"""
|
|
|
- Extracts the environment variables from a source-able script on *unix.
|
|
|
- Needs to be sourced in a bash shell.
|
|
|
+ Extracts the environment variables from a source-able script on Unix.
|
|
|
+ On Windows, uses PowerShell to source the script and capture env changes.
|
|
|
"""
|
|
|
- # Assume the script exists and is executable
|
|
|
for script_path in script_paths:
|
|
|
if not script_path.is_file():
|
|
|
raise Exception(
|
|
|
f"The file '{script_path}' does not exist or is not a file."
|
|
|
)
|
|
|
|
|
|
- # Parse the result output of executing "env"
|
|
|
+ def parse_env(env_str):
|
|
|
+ env = {}
|
|
|
+ for line in env_str.splitlines():
|
|
|
+ if '=' in line:
|
|
|
+ key, value = line.split('=', 1)
|
|
|
+ env[key] = value
|
|
|
+ return env
|
|
|
+
|
|
|
+ system = platform.system().lower()
|
|
|
+
|
|
|
+ if system == "windows":
|
|
|
+ return _extract_env_via_powershell(script_paths)
|
|
|
+ else:
|
|
|
+ return _extract_env_via_bash(script_paths)
|
|
|
+
|
|
|
+
|
|
|
+def _extract_env_via_bash(script_paths: List[Path]) -> Dict[str, str]:
|
|
|
def parse_env(env_str):
|
|
|
env = {}
|
|
|
for line in env_str.splitlines():
|
|
|
@@ -26,7 +42,6 @@ def extract_unix_vars_of_source(script_paths: List[Path]) -> Dict[str, str]:
|
|
|
return env
|
|
|
|
|
|
try:
|
|
|
- # Get original environment variables
|
|
|
original_env_output = subprocess.check_output(
|
|
|
['bash', '-c', 'env'],
|
|
|
stderr=subprocess.PIPE,
|
|
|
@@ -34,12 +49,10 @@ def extract_unix_vars_of_source(script_paths: List[Path]) -> Dict[str, str]:
|
|
|
)
|
|
|
original = parse_env(original_env_output)
|
|
|
|
|
|
- # Merge all sourcing script paths in to one command
|
|
|
source_command = ' && '.join(
|
|
|
[f'source {script_path}' for script_path in script_paths]
|
|
|
)
|
|
|
|
|
|
- # Get the environment variables after sourcing the script
|
|
|
sourced_env_output = subprocess.check_output(
|
|
|
['bash', '-c', f'{source_command} && env'],
|
|
|
stderr=subprocess.PIPE,
|
|
|
@@ -47,7 +60,6 @@ def extract_unix_vars_of_source(script_paths: List[Path]) -> Dict[str, str]:
|
|
|
)
|
|
|
sourced = parse_env(sourced_env_output)
|
|
|
|
|
|
- # Get the difference
|
|
|
diff = {
|
|
|
k: v
|
|
|
for k, v in sourced.items()
|
|
|
@@ -61,6 +73,56 @@ def extract_unix_vars_of_source(script_paths: List[Path]) -> Dict[str, str]:
|
|
|
)
|
|
|
|
|
|
|
|
|
+def _extract_env_via_powershell(script_paths: List[Path]) -> Dict[str, str]:
|
|
|
+ """
|
|
|
+ Uses PowerShell to source a script (e.g. .ps1 or env-setup script)
|
|
|
+ and capture environment variable changes.
|
|
|
+ """
|
|
|
+ try:
|
|
|
+ # Build a PowerShell command that captures env before and after sourcing
|
|
|
+ script_args = ' '.join(
|
|
|
+ [f'"{str(p)}"' for p in script_paths]
|
|
|
+ )
|
|
|
+ ps_command = (
|
|
|
+ '$before = Get-ChildItem Env: | ForEach-Object { "$($_.Name)=$($_.Value)" }; '
|
|
|
+ f'{script_args}; '
|
|
|
+ '$after = Get-ChildItem Env: | ForEach-Object { "$($_.Name)=$($_.Value)" }; '
|
|
|
+ 'Write-Output "---BEFORE---"; '
|
|
|
+ 'Write-Output $before; '
|
|
|
+ 'Write-Output "---AFTER---"; '
|
|
|
+ 'Write-Output $after'
|
|
|
+ )
|
|
|
+
|
|
|
+ output = subprocess.check_output(
|
|
|
+ ['powershell', '-NoProfile', '-NonInteractive', '-Command', ps_command],
|
|
|
+ stderr=subprocess.PIPE,
|
|
|
+ text=True,
|
|
|
+ )
|
|
|
+
|
|
|
+ # Split output into before and after sections
|
|
|
+ parts = output.split('---BEFORE---')
|
|
|
+ if len(parts) < 2:
|
|
|
+ return {}
|
|
|
+ after_section = parts[1].split('---AFTER---')
|
|
|
+ before_str = after_section[0]
|
|
|
+ after_str = after_section[1] if len(after_section) > 1 else ''
|
|
|
+
|
|
|
+ before = parse_env(before_str)
|
|
|
+ after = parse_env(after_str)
|
|
|
+
|
|
|
+ diff = {
|
|
|
+ k: v
|
|
|
+ for k, v in after.items()
|
|
|
+ if k not in before or before.get(k) != v
|
|
|
+ }
|
|
|
+
|
|
|
+ return diff
|
|
|
+ except subprocess.CalledProcessError as e:
|
|
|
+ raise Exception(
|
|
|
+ f"Failed to extract environment variables from [{script_paths}] via PowerShell: {e.stderr}"
|
|
|
+ )
|
|
|
+
|
|
|
+
|
|
|
def get_gpustack_env(env_var: str) -> Optional[str]:
|
|
|
env_name = "GPUSTACK_" + env_var
|
|
|
return os.getenv(env_name)
|