image-match/python/py/Lib/site-packages/lightning_utilities/install/requirements.py

# Licensed under the Apache License, Version 2.0 (the "License");
#     http://www.apache.org/licenses/LICENSE-2.0
#
"""Utilities to parse and adjust Python requirements files.

This module parses requirement lines while preserving inline comments and pip arguments and
supports relaxing version pins based on a chosen unfreeze strategy: "none", "major", or "all".

"""

import re
from collections.abc import Iterable, Iterator
from pathlib import Path
from typing import Any, Optional, Union

from packaging.requirements import Requirement
from packaging.version import Version


def _yield_lines(strs: Union[str, Iterable[str]]) -> Iterator[str]:
    """Yield non-empty, non-comment lines from a string or iterable of strings.

    Adapted from pkg_resources.yield_lines.

    """
    if isinstance(strs, str):
        strs = strs.splitlines()
    for line in strs:
        line = line.strip()
        if line and not line.startswith("#"):
            yield line


class _RequirementWithComment(Requirement):
    """Requirement subclass that preserves an inline comment and optional pip argument.

    Attributes:
        comment: The trailing comment captured from the requirement line (including the leading '# ...').
        pip_argument: A preceding pip argument line (e.g., ``"--extra-index-url ..."``) associated
            with this requirement, or ``None`` if not provided.
        strict: Whether the special marker ``"# strict"`` appears in ``comment`` (case-insensitive), in which case
            upper bound adjustments are disabled.

    """

    strict_string = "# strict"

    def __init__(self, *args: Any, comment: str = "", pip_argument: Optional[str] = None, **kwargs: Any) -> None:
        super().__init__(*args, **kwargs)
        self.comment = comment
        if not (pip_argument is None or pip_argument):  # sanity check that it's not an empty str
            raise RuntimeError(f"wrong pip argument: {pip_argument}")
        self.pip_argument = pip_argument
        self.strict = self.strict_string in comment.lower()

    def adjust(self, unfreeze: str) -> str:
        """Adjust version specifiers according to the selected unfreeze strategy.

        The special marker ``"# strict"`` in the captured comment disables any relaxation of upper bounds.

        >>> _RequirementWithComment("arrow<=1.2.2,>=1.2.0", comment="# anything").adjust("none")
        'arrow<=1.2.2,>=1.2.0'
        >>> _RequirementWithComment("arrow<=1.2.2,>=1.2.0", comment="# strict").adjust("none")
        'arrow<=1.2.2,>=1.2.0  # strict'
        >>> _RequirementWithComment("arrow<=1.2.2,>=1.2.0", comment="# my name").adjust("all")
        'arrow>=1.2.0'
        >>> _RequirementWithComment("arrow>=1.2.0, <=1.2.2", comment="# strict").adjust("all")
        'arrow<=1.2.2,>=1.2.0  # strict'
        >>> _RequirementWithComment("arrow").adjust("all")
        'arrow'
        >>> _RequirementWithComment("arrow>=1.2.0, <=1.2.2", comment="# cool").adjust("major")
        'arrow<2.0,>=1.2.0'
        >>> _RequirementWithComment("arrow>=1.2.0, <=1.2.2", comment="# strict").adjust("major")
        'arrow<=1.2.2,>=1.2.0  # strict'
        >>> _RequirementWithComment("arrow>=1.2.0").adjust("major")
        'arrow>=1.2.0'
        >>> _RequirementWithComment("arrow").adjust("major")
        'arrow'

        Args:
            unfreeze: One of:
                - ``"none"``: Keep all version specifiers unchanged.
                - ``"major"``: Relax the upper bound to the next major version (e.g., ``<2.0``).
                - ``"all"``: Drop any upper bound constraint entirely.

        Returns:
            The adjusted requirement string. If strict, the original string is returned with the strict marker appended.

        """
        out = str(self)
        if self.strict:
            return f"{out}  {self.strict_string}"
        if unfreeze == "major":
            for spec in self.specifier:
                if spec.operator in ("<", "<="):
                    major = Version(spec.version).major
                    # replace upper bound with major version increased by one
                    return out.replace(f"{spec.operator}{spec.version}", f"<{int(major) + 1}.0")
        elif unfreeze == "all":
            for spec in self.specifier:
                if spec.operator in ("<", "<="):
                    # drop upper bound (with or without trailing/leading comma)
                    upper = f"{spec.operator}{spec.version}"
                    result = out.replace(f"{upper},", "").replace(f",{upper}", "")
                    if upper in result:
                        result = result.replace(upper, "")
                    return result.strip()
        elif unfreeze != "none":
            raise ValueError(f"Unexpected unfreeze: {unfreeze!r} value.")
        return out


def _parse_requirements(strs: Union[str, Iterable[str]]) -> Iterator[_RequirementWithComment]:
    r"""Adapted from ``pkg_resources.parse_requirements`` to include comments and pip arguments.

    Parses a sequence or string of requirement lines, preserving trailing comments and associating any
    preceding pip arguments (``--...``) with the subsequent requirement. Lines starting with ``-r`` or
    containing direct URLs are ignored.

    >>> txt = ['# ignored', '', 'this # is an', '--piparg', 'example', 'foo # strict', 'thing', '-r different/file.txt']
    >>> [r.adjust('none') for r in _parse_requirements(txt)]
    ['this', 'example', 'foo  # strict', 'thing']
    >>> txt = '\\n'.join(txt)
    >>> [r.adjust('none') for r in _parse_requirements(txt)]
    ['this', 'example', 'foo  # strict', 'thing']

    Args:
        strs: Either an iterable of requirement lines or a single multi-line string.

    Yields:
        _RequirementWithComment: Parsed requirement objects with preserved comment and pip argument.

    """
    lines = _yield_lines(strs)
    pip_argument = None
    for line in lines:
        # Drop comments -- a hash without a space may be in a URL.
        if " #" in line:
            comment_pos = line.find(" #")
            line, comment = line[:comment_pos], line[comment_pos:]
        else:
            comment = ""
        # If there is a line continuation, drop it, and append the next line.
        if line.endswith("\\"):
            line = line[:-1].strip()
            try:
                line += next(lines)
            except StopIteration:
                return
        # If there's a pip argument, save it
        if line.startswith("--"):
            pip_argument = line
            continue
        if line.startswith("-r "):
            # linked requirement files are unsupported
            continue
        if "@" in line or re.search("https?://", line):
            # skip lines with links like `pesq @ git+https://github.com/ludlows/python-pesq`
            continue
        yield _RequirementWithComment(line, comment=comment, pip_argument=pip_argument)
        pip_argument = None


def load_requirements(path_dir: str, file_name: str = "base.txt", unfreeze: str = "all") -> list[str]:
    """Load, parse, and optionally relax requirement specifiers from a file.

    >>> import os
    >>> from lightning_utilities import _PROJECT_ROOT
    >>> path_req = os.path.join(_PROJECT_ROOT, "requirements")
    >>> load_requirements(path_req, "docs.txt", unfreeze="major")  # doctest: +ELLIPSIS +NORMALIZE_WHITESPACE
    ['sphinx<6.0,>=4.0', ...]

    Args:
        path_dir: Directory containing the requirements file.
        file_name: The requirements filename inside ``path_dir``.
        unfreeze: Unfreeze strategy: ``"none"``, ``"major"``, or ``"all"`` (see ``_RequirementWithComment.adjust``).

    Returns:
        A list of requirement strings adjusted according to ``unfreeze``.

    Raises:
        ValueError: If ``unfreeze`` is not one of the supported options.
        FileNotFoundError: If the composed path does not exist.

    """
    if unfreeze not in {"none", "major", "all"}:
        raise ValueError(f'unsupported option of "{unfreeze}"')
    path = Path(path_dir) / file_name
    if not path.exists():
        raise FileNotFoundError(f"missing file for {(path_dir, file_name, path)}")
    text = path.read_text()
    return [req.adjust(unfreeze) for req in _parse_requirements(text)]