276 lines
7.6 KiB
Python
276 lines
7.6 KiB
Python
"""Shared helpers for the AI Toolkit manager. Stdlib only."""
|
|
|
|
import json
|
|
import os
|
|
import platform
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
|
|
REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
|
|
IS_WINDOWS = platform.system() == "Windows"
|
|
IS_MAC = platform.system() == "Darwin"
|
|
IS_LINUX = platform.system() == "Linux"
|
|
|
|
USER_AGENT = (
|
|
"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 "
|
|
"(KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36"
|
|
)
|
|
|
|
# When --json is used, human output goes to stderr so stdout stays machine-readable
|
|
_json_mode = False
|
|
|
|
|
|
def set_json_mode(enabled):
|
|
global _json_mode
|
|
_json_mode = enabled
|
|
|
|
|
|
def _supports_color(stream):
|
|
if os.environ.get("NO_COLOR"):
|
|
return False
|
|
return hasattr(stream, "isatty") and stream.isatty()
|
|
|
|
|
|
def _emit(prefix, msg, color):
|
|
stream = sys.stderr if _json_mode else sys.stdout
|
|
if _supports_color(stream):
|
|
stream.write("\033[%sm%s\033[0m %s\n" % (color, prefix, msg))
|
|
else:
|
|
stream.write("%s %s\n" % (prefix, msg))
|
|
stream.flush()
|
|
|
|
|
|
def info(msg):
|
|
_emit("[*]", msg, "36")
|
|
|
|
|
|
def ok(msg):
|
|
_emit("[+]", msg, "32")
|
|
|
|
|
|
def warn(msg):
|
|
_emit("[!]", msg, "33")
|
|
|
|
|
|
def error(msg):
|
|
_emit("[x]", msg, "31")
|
|
|
|
|
|
def die(msg, code=1):
|
|
error(msg)
|
|
sys.exit(code)
|
|
|
|
|
|
def print_json(data):
|
|
sys.stdout.write(json.dumps(data, indent=2) + "\n")
|
|
sys.stdout.flush()
|
|
|
|
|
|
def run(cmd, cwd=None, capture=False, check=True, env=None, stream=False):
|
|
"""Run a command. capture=True returns stdout text (stripped).
|
|
|
|
stream=True inherits stdio so the user sees live output.
|
|
Returns (returncode, stdout_or_None).
|
|
"""
|
|
kwargs = {"cwd": cwd or REPO_ROOT}
|
|
if env is not None:
|
|
kwargs["env"] = env
|
|
if capture:
|
|
kwargs["stdout"] = subprocess.PIPE
|
|
kwargs["stderr"] = subprocess.PIPE
|
|
elif _json_mode and not stream:
|
|
# keep stdout clean in json mode
|
|
kwargs["stdout"] = sys.stderr
|
|
|
|
try:
|
|
proc = subprocess.run(cmd, **kwargs)
|
|
except FileNotFoundError:
|
|
if check:
|
|
die("Command not found: %s" % cmd[0])
|
|
return 127, None
|
|
|
|
out = None
|
|
if capture:
|
|
out = proc.stdout.decode("utf-8", errors="replace").strip()
|
|
if check and proc.returncode != 0:
|
|
detail = ""
|
|
if capture and proc.stderr:
|
|
detail = "\n" + proc.stderr.decode("utf-8", errors="replace").strip()
|
|
die("Command failed (%d): %s%s" % (proc.returncode, " ".join(cmd), detail))
|
|
return proc.returncode, out
|
|
|
|
|
|
def which(name):
|
|
return shutil.which(name)
|
|
|
|
|
|
def find_uv():
|
|
"""Find uv: repo-local .uv/ first, then PATH, then common install dirs."""
|
|
home = os.path.expanduser("~")
|
|
candidates = [
|
|
os.path.join(REPO_ROOT, ".uv", "uv.exe" if IS_WINDOWS else "uv"),
|
|
]
|
|
for c in candidates:
|
|
if os.path.isfile(c) and os.access(c, os.X_OK):
|
|
return c
|
|
uv = shutil.which("uv")
|
|
if uv:
|
|
return uv
|
|
candidates = [
|
|
os.path.join(home, ".local", "bin", "uv"),
|
|
os.path.join(home, ".cargo", "bin", "uv"),
|
|
]
|
|
if IS_WINDOWS:
|
|
local = os.environ.get("LOCALAPPDATA", "")
|
|
if local:
|
|
candidates.append(os.path.join(local, "uv", "uv.exe"))
|
|
candidates.append(os.path.join(home, ".local", "bin", "uv.exe"))
|
|
for c in candidates:
|
|
if os.path.isfile(c) and os.access(c, os.X_OK):
|
|
return c
|
|
return None
|
|
|
|
|
|
def venv_dir():
|
|
"""Existing venv dir (.venv preferred, matching ui/cron/pythonPath.ts), else default target."""
|
|
for name in (".venv", "venv"):
|
|
d = os.path.join(REPO_ROOT, name)
|
|
if os.path.isdir(d):
|
|
return d
|
|
return os.path.join(REPO_ROOT, ".venv")
|
|
|
|
|
|
def venv_python(venv=None):
|
|
venv = venv or venv_dir()
|
|
if IS_WINDOWS:
|
|
return os.path.join(venv, "Scripts", "python.exe")
|
|
return os.path.join(venv, "bin", "python3")
|
|
|
|
|
|
# Env vars that let a system/conda/pyenv Python leak into our subprocesses.
|
|
# Scrubbed from every python/pip/node invocation (mirrors what the community
|
|
# Windows installer learned the hard way).
|
|
_SCRUB_VARS = (
|
|
"PYTHONPATH",
|
|
"PYTHONHOME",
|
|
"PYTHON",
|
|
"PYTHONSTARTUP",
|
|
"PYTHONUSERBASE",
|
|
"PYTHONEXECUTABLE",
|
|
"PIP_CONFIG_FILE",
|
|
"PIP_REQUIRE_VIRTUALENV",
|
|
"VIRTUAL_ENV",
|
|
"CONDA_PREFIX",
|
|
"CONDA_DEFAULT_ENV",
|
|
"PYENV_ROOT",
|
|
"PYENV_VERSION",
|
|
)
|
|
|
|
|
|
def clean_env(extra=None):
|
|
"""os.environ copy with Python-hijacking vars removed.
|
|
|
|
Also points uv's managed-python store into the repo (.uv/python) so
|
|
interpreter downloads never land outside the checkout.
|
|
"""
|
|
env = os.environ.copy()
|
|
for var in _SCRUB_VARS:
|
|
env.pop(var, None)
|
|
env.setdefault("UV_PYTHON_INSTALL_DIR", os.path.join(REPO_ROOT, ".uv", "python"))
|
|
if extra:
|
|
env.update(extra)
|
|
return env
|
|
|
|
|
|
def download(url, dest, label=None):
|
|
"""Download url to dest (stdlib only), logging progress every ~10%.
|
|
|
|
Sends a browser User-Agent: some mirrors (ffmpeg.martin-riedl.de) return
|
|
403 for the default "Python-urllib/x.y" agent.
|
|
"""
|
|
import urllib.request
|
|
|
|
label = label or os.path.basename(dest)
|
|
info("Downloading %s ..." % label)
|
|
last = [-1]
|
|
|
|
def report(read, total):
|
|
if total <= 0:
|
|
return
|
|
pct = min(100, int(read * 100 / total))
|
|
if pct >= last[0] + 10:
|
|
last[0] = pct
|
|
info(" %s: %d%%" % (label, pct))
|
|
|
|
req = urllib.request.Request(url, headers={"User-Agent": USER_AGENT})
|
|
tmp = dest + ".part"
|
|
try:
|
|
with urllib.request.urlopen(req) as resp, open(tmp, "wb") as out:
|
|
total = int(resp.headers.get("Content-Length") or 0)
|
|
read = 0
|
|
while True:
|
|
chunk = resp.read(1 << 16)
|
|
if not chunk:
|
|
break
|
|
out.write(chunk)
|
|
read += len(chunk)
|
|
report(read, total)
|
|
except Exception as e: # noqa: BLE001 - surface any network failure clearly
|
|
if os.path.exists(tmp):
|
|
os.remove(tmp)
|
|
# Windows: python's ssl verifies against the OS cert store but never
|
|
# triggers Windows' on-demand intermediate-CA fetching, so on a fresh
|
|
# machine github downloads can fail CERTIFICATE_VERIFY_FAILED even
|
|
# though the chain is fine. curl (bundled since Win10, schannel-based)
|
|
# does fetch intermediates — fall back to it before giving up.
|
|
if not _download_with_curl(url, tmp, label):
|
|
die("Download failed for %s: %s" % (url, e))
|
|
os.replace(tmp, dest)
|
|
|
|
|
|
def _download_with_curl(url, tmp, label):
|
|
curl = shutil.which("curl")
|
|
if not curl:
|
|
return False
|
|
info(" %s: retrying with curl..." % label)
|
|
try:
|
|
code = subprocess.call(
|
|
[curl, "-fSL", "--retry", "3", "-o", tmp, url],
|
|
stdout=subprocess.DEVNULL,
|
|
stderr=subprocess.DEVNULL,
|
|
)
|
|
except OSError:
|
|
return False
|
|
if code != 0 or not os.path.exists(tmp):
|
|
if os.path.exists(tmp):
|
|
os.remove(tmp)
|
|
return False
|
|
return True
|
|
|
|
|
|
def extract_archive(archive, dest_dir):
|
|
"""Extract .zip / .tar.* into dest_dir (created if needed)."""
|
|
import tarfile
|
|
import zipfile
|
|
|
|
os.makedirs(dest_dir, exist_ok=True)
|
|
if archive.endswith(".zip"):
|
|
with zipfile.ZipFile(archive) as z:
|
|
z.extractall(dest_dir)
|
|
else:
|
|
with tarfile.open(archive) as t:
|
|
t.extractall(dest_dir)
|
|
|
|
|
|
def file_hash(paths):
|
|
import hashlib
|
|
|
|
h = hashlib.sha256()
|
|
for p in sorted(paths):
|
|
if os.path.isfile(p):
|
|
with open(p, "rb") as f:
|
|
h.update(f.read())
|
|
return h.hexdigest()
|