AstrBot/astrbot/core/utils/requirements_utils.py

import importlib.metadata as importlib_metadata
import logging
import os
import re
import shlex
import sys
from collections.abc import Iterable, Iterator
from dataclasses import dataclass

from packaging.requirements import InvalidRequirement, Requirement
from packaging.specifiers import SpecifierSet
from packaging.version import InvalidVersion, Version

from astrbot.core.utils.astrbot_path import get_astrbot_site_packages_path
from astrbot.core.utils.runtime_env import is_packaged_desktop_runtime

logger = logging.getLogger("astrbot")


class RequirementsPrecheckFailed(Exception):
    """Raised when the pre-check of requirements fails."""

    pass


@dataclass(frozen=True)
class ParsedPackageInput:
    specs: tuple[str, ...]
    requirement_names: frozenset[str]


def canonicalize_distribution_name(name: str) -> str:
    return re.sub(r"[-_.]+", "-", name).strip("-").lower()


def strip_inline_requirement_comment(raw_input: str) -> str:
    if raw_input.lstrip().startswith("#"):
        return ""
    return re.split(r"[ \t]+#", raw_input, maxsplit=1)[0].strip()


def _specifier_contains_version(specifier: SpecifierSet, version: str) -> bool:
    try:
        parsed_version = Version(version)
    except InvalidVersion:
        return False
    return specifier.contains(parsed_version, prereleases=True)


def _looks_like_local_path_reference(token: str) -> bool:
    candidate = token.strip()
    if not candidate:
        return False
    return candidate in {".", ".."} or candidate.startswith(
        ("./", "../", "/", "~/", ".\\", "..\\", "\\")
    )


def looks_like_direct_reference(token: str) -> bool:
    candidate = token.strip()
    if not candidate:
        return False
    return (
        _looks_like_local_path_reference(candidate)
        or candidate.startswith("git+")
        or "://" in candidate
    )


def extract_requirement_name(raw_requirement: str) -> str | None:
    line = raw_requirement.split("#", 1)[0].strip()
    if not line:
        return None
    if line.startswith(("-r", "--requirement", "-c", "--constraint")):
        return None

    egg_match = re.search(r"#egg=([A-Za-z0-9_.-]+)", raw_requirement)
    if egg_match:
        return canonicalize_distribution_name(egg_match.group(1))

    if line.startswith("-"):
        return None

    candidate = re.split(r"[<>=!~;\s\[]", line, maxsplit=1)[0].strip()
    if not candidate:
        return None
    return canonicalize_distribution_name(candidate)


def _parse_editable_or_direct_name(target: str) -> str | None:
    name = extract_requirement_name(target)
    if not name:
        return None
    if "#egg=" in target or not looks_like_direct_reference(target):
        return name
    return None


def _parse_requirement_name_and_spec(
    line: str,
) -> tuple[str | None, SpecifierSet | None]:
    if line.startswith(("-c", "--constraint")):
        return None, None

    try:
        req = Requirement(line)
    except InvalidRequirement:
        tokens = shlex.split(line)
        if not tokens:
            return None, None

        editable_target: str | None = None
        if tokens[0] in {"-e", "--editable"} and len(tokens) > 1:
            editable_target = tokens[1]
        elif tokens[0].startswith("--editable="):
            editable_target = tokens[0].split("=", 1)[1]

        if editable_target:
            name = _parse_editable_or_direct_name(editable_target)
            return (name, None) if name else (None, None)

        name = _parse_editable_or_direct_name(line)
        return (name, None) if name else (None, None)

    if req.marker and not req.marker.evaluate():
        return None, None

    return canonicalize_distribution_name(req.name), (req.specifier or None)


def _parse_requirement_line(
    line: str,
) -> tuple[str, SpecifierSet | None] | None:
    name, specifier = _parse_requirement_name_and_spec(line)
    return (name, specifier) if name else None


def _extract_requirement_names_from_package_tokens(tokens: list[str]) -> frozenset[str]:
    requirement_names: set[str] = set()
    skip_next_for: str | None = None

    for token in tokens:
        if skip_next_for:
            if skip_next_for == "editable":
                name = _parse_editable_or_direct_name(token)
                if name:
                    requirement_names.add(name)
            skip_next_for = None
            continue

        if token in {"-e", "--editable"}:
            skip_next_for = "editable"
            continue

        if token in {
            "-i",
            "--index-url",
            "--extra-index-url",
            "-f",
            "--find-links",
            "--trusted-host",
            "-r",
            "--requirement",
            "-c",
            "--constraint",
        }:
            skip_next_for = "option-value"
            continue

        if token.startswith(("--editable=",)):
            editable_target = token.split("=", 1)[1]
            name = _parse_editable_or_direct_name(editable_target)
            if name:
                requirement_names.add(name)
            continue

        if token.startswith(
            (
                "--index-url=",
                "--extra-index-url=",
                "--find-links=",
                "--trusted-host=",
                "--requirement=",
                "--constraint=",
            )
        ):
            continue

        if (
            (token.startswith("-i") and token != "-i")
            or (token.startswith("-f") and token != "-f")
            or token == "--no-index"
        ):
            continue

        if token.startswith("-"):
            continue

        name, _ = _parse_requirement_name_and_spec(token)
        if name:
            requirement_names.add(name)

    return frozenset(requirement_names)


def parse_package_install_input(raw_input: str) -> ParsedPackageInput:
    specs: list[str] = []
    requirement_names: set[str] = set()
    normalized = raw_input.strip()
    if not normalized:
        return ParsedPackageInput(specs=(), requirement_names=frozenset())

    for raw_line in normalized.splitlines():
        line = strip_inline_requirement_comment(raw_line)
        if not line:
            continue

        try:
            Requirement(line)
        except InvalidRequirement:
            tokens = shlex.split(line)
            if not tokens:
                continue
            specs.extend(tokens)
            requirement_names.update(
                _extract_requirement_names_from_package_tokens(tokens)
            )
            continue

        specs.append(line)
        name, _ = _parse_requirement_name_and_spec(line)
        if name:
            requirement_names.add(name)

    return ParsedPackageInput(
        specs=tuple(specs),
        requirement_names=frozenset(requirement_names),
    )


def _iter_requirement_lines(
    requirements_path: str,
    _visited: set[str] | None = None,
) -> Iterator[str]:
    visited = _visited or set()
    resolved_path = os.path.realpath(requirements_path)
    if resolved_path in visited:
        logger.warning(
            "检测到循环依赖的 requirements 包含: %s，将跳过该文件", resolved_path
        )
        return
    visited.add(resolved_path)

    with open(resolved_path, encoding="utf-8") as f:
        for raw_line in f:
            line = strip_inline_requirement_comment(raw_line)
            if not line:
                continue

            tokens = shlex.split(line)
            if not tokens:
                continue

            nested: str | None = None
            if tokens[0] in {"-r", "--requirement"} and len(tokens) > 1:
                nested = tokens[1]
            elif tokens[0].startswith("--requirement="):
                nested = tokens[0].split("=", 1)[1]

            if nested:
                if not os.path.isabs(nested):
                    nested = os.path.join(os.path.dirname(resolved_path), nested)
                yield from _iter_requirement_lines(nested, _visited=visited)
                continue

            yield line


def iter_requirements(
    requirements_path: str | None = None,
    lines: Iterable[str] | None = None,
) -> Iterator[tuple[str, SpecifierSet | None]]:
    if lines is None:
        if requirements_path is None:
            raise ValueError("Either requirements_path or lines must be provided")
        lines = _iter_requirement_lines(requirements_path)

    for line in lines:
        parsed = _parse_requirement_line(line)
        if parsed is not None:
            yield parsed


def extract_requirement_names(requirements_path: str) -> set[str]:
    try:
        return {
            name for name, _ in iter_requirements(requirements_path=requirements_path)
        }
    except Exception as exc:
        logger.warning("读取依赖文件失败，跳过冲突检测: %s", exc)
        return set()


def get_requirement_check_paths() -> list[str]:
    paths = list(sys.path)
    if is_packaged_desktop_runtime():
        target_site_packages = get_astrbot_site_packages_path()
        if os.path.isdir(target_site_packages):
            paths.insert(0, target_site_packages)
    return paths


def _canonical_distribution_identity(distribution) -> tuple[str | None, str | None]:
    distribution_name = (
        distribution.metadata["Name"] if "Name" in distribution.metadata else None
    )
    if not distribution_name:
        return None, None
    return canonicalize_distribution_name(distribution_name), distribution.version


def collect_installed_distribution_versions(paths: list[str]) -> dict[str, str] | None:
    installed: dict[str, str] = {}
    try:
        for distribution in importlib_metadata.distributions(path=paths):
            distribution_name, version = _canonical_distribution_identity(distribution)
            if not distribution_name or not version:
                continue
            installed.setdefault(distribution_name, version)
    except Exception as exc:
        logger.warning("读取已安装依赖失败，跳过缺失依赖预检查: %s", exc)
        return None
    return installed


def _load_requirement_lines_for_precheck(
    requirements_path: str,
) -> tuple[bool, list[str] | None]:
    try:
        requirement_lines = list(_iter_requirement_lines(requirements_path))
    except Exception as exc:
        logger.warning(
            "预检查缺失依赖失败，将回退到完整安装: %s (%s)",
            requirements_path,
            exc,
        )
        return False, None

    fallback_line = next(
        (
            line
            for line in requirement_lines
            if (
                (
                    line.startswith(("-e ", "--editable ", "--editable="))
                    and "#egg=" not in line
                )
                or (
                    _parse_requirement_line(line) is None
                    and looks_like_direct_reference(line)
                )
            )
        ),
        None,
    )
    if fallback_line is not None:
        logger.warning(
            "预检查缺失依赖失败，将回退到完整安装: unresolved direct reference in %s: %s",
            requirements_path,
            fallback_line,
        )
        return False, None

    return True, requirement_lines


def find_missing_requirements(requirements_path: str) -> set[str] | None:
    can_precheck, requirement_lines = _load_requirement_lines_for_precheck(
        requirements_path
    )
    if not can_precheck or requirement_lines is None:
        return None

    required = list(iter_requirements(lines=requirement_lines))
    if not required:
        return set()

    installed = collect_installed_distribution_versions(get_requirement_check_paths())
    if installed is None:
        return None

    missing: set[str] = set()
    for name, specifier in required:
        installed_version = installed.get(name)
        if not installed_version:
            missing.add(name)
            continue
        if specifier and not _specifier_contains_version(specifier, installed_version):
            missing.add(name)

    return missing


def find_missing_requirements_or_raise(requirements_path: str) -> set[str]:
    missing = find_missing_requirements(requirements_path)
    if missing is None:
        raise RequirementsPrecheckFailed(f"预检查失败: {requirements_path}")
    return missing