parse_requirements.bzl

# Copyright 2024 The Bazel Authors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Requirements parsing for whl_library creation.

Use cases that the code needs to cover:
* A single requirements_lock file that is used for the host platform.
* Per-OS requirements_lock files that are used for the host platform.
* A target platform specific requirements_lock that is used with extra
  pip arguments with --platform, etc and download_only = True.

In the last case only a single `requirements_lock` file is allowed, in all
other cases we assume that there may be a desire to resolve the requirements
file for the host platform to be backwards compatible with the legacy
behavior.
"""

load("//python/private:normalize_name.bzl", "normalize_name")
load("//python/private:repo_utils.bzl", "repo_utils")
load(":index_sources.bzl", "index_sources")
load(":parse_requirements_txt.bzl", "parse_requirements_txt")
load(":whl_target_platforms.bzl", "select_whls")

def _extract_version(entry):
    """Extract the version part from the requirement string.


    Args:
        entry: {type}`str` The requirement string.
    """
    version_start = entry.find("==")
    if version_start != -1:
        # Extract everything after '==' until the next space or end of the string
        version, _, _ = entry[version_start + 2:].partition(" ")
        return version
    return None

def parse_requirements(
        ctx,
        *,
        requirements_by_platform = {},
        extra_pip_args = [],
        get_index_urls = None,
        evaluate_markers = None,
        logger = None):
    """Get the requirements with platforms that the requirements apply to.

    Args:
        ctx: A context that has .read function that would read contents from a label.
        requirements_by_platform (label_keyed_string_dict): a way to have
            different package versions (or different packages) for different
            os, arch combinations.
        extra_pip_args (string list): Extra pip arguments to perform extra validations and to
            be joined with args found in files.
        get_index_urls: Callable[[ctx, list[str]], dict], a callable to get all
            of the distribution URLs from a PyPI index. Accepts ctx and
            distribution names to query.
        evaluate_markers: A function to use to evaluate the requirements.
            Accepts a dict where keys are requirement lines to evaluate against
            the platforms stored as values in the input dict. Returns the same
            dict, but with values being platforms that are compatible with the
            requirements line.
        logger: repo_utils.logger or None, a simple struct to log diagnostic messages.

    Returns:
        {type}`dict[str, list[struct]]` where the key is the distribution name and the struct
        contains the following attributes:
         * `distribution`: {type}`str` The non-normalized distribution name.
         * `srcs`: {type}`struct` The parsed requirement line for easier Simple
           API downloading (see `index_sources` return value).
         * `target_platforms`: {type}`list[str]` Target platforms that this package is for.
             The format is `cp3{minor}_{os}_{arch}`.
         * `is_exposed`: {type}`bool` `True` if the package should be exposed via the hub
           repository.
         * `extra_pip_args`: {type}`list[str]` pip args to use in case we are
           not using the bazel downloader to download the archives. This should
           be passed to {obj}`whl_library`.
         * `whls`: {type}`list[struct]` The list of whl entries that can be
           downloaded using the bazel downloader.
         * `sdist`: {type}`list[struct]` The sdist that can be downloaded using
           the bazel downloader.

        The second element is extra_pip_args should be passed to `whl_library`.
    """
    evaluate_markers = evaluate_markers or (lambda _: {})
    options = {}
    requirements = {}
    for file, plats in requirements_by_platform.items():
        if logger:
            logger.debug(lambda: "Using {} for {}".format(file, plats))
        contents = ctx.read(file)

        # Parse the requirements file directly in starlark to get the information
        # needed for the whl_library declarations later.
        parse_result = parse_requirements_txt(contents)

        # Replicate a surprising behavior that WORKSPACE builds allowed:
        # Defining a repo with the same name multiple times, but only the last
        # definition is respected.
        # The requirement lines might have duplicate names because lines for extras
        # are returned as just the base package name. e.g., `foo[bar]` results
        # in an entry like `("foo", "foo[bar] == 1.0 ...")`.
        requirements_dict = {
            (normalize_name(entry[0]), _extract_version(entry[1])): entry
            for entry in sorted(
                parse_result.requirements,
                # Get the longest match and fallback to original WORKSPACE sorting,
                # which should get us the entry with most extras.
                #
                # FIXME @aignas 2024-05-13: The correct behaviour might be to get an
                # entry with all aggregated extras, but it is unclear if we
                # should do this now.
                key = lambda x: (len(x[1].partition("==")[0]), x),
            )
        }.values()

        tokenized_options = []
        for opt in parse_result.options:
            for p in opt.split(" "):
                tokenized_options.append(p)

        pip_args = tokenized_options + extra_pip_args
        for plat in plats:
            requirements[plat] = requirements_dict
            options[plat] = pip_args

    requirements_by_platform = {}
    reqs_with_env_markers = {}
    for target_platform, reqs_ in requirements.items():
        extra_pip_args = options[target_platform]

        for distribution, requirement_line in reqs_:
            for_whl = requirements_by_platform.setdefault(
                normalize_name(distribution),
                {},
            )

            if ";" in requirement_line:
                reqs_with_env_markers.setdefault(requirement_line, []).append(target_platform)

            for_req = for_whl.setdefault(
                (requirement_line, ",".join(extra_pip_args)),
                struct(
                    distribution = distribution,
                    srcs = index_sources(requirement_line),
                    requirement_line = requirement_line,
                    target_platforms = [],
                    extra_pip_args = extra_pip_args,
                ),
            )
            for_req.target_platforms.append(target_platform)

    # This may call to Python, so execute it early (before calling to the
    # internet below) and ensure that we call it only once.
    #
    # NOTE @aignas 2024-07-13: in the future, if this is something that we want
    # to do, we could use Python to parse the requirement lines and infer the
    # URL of the files to download things from. This should be important for
    # VCS package references.
    env_marker_target_platforms = evaluate_markers(reqs_with_env_markers)
    if logger:
        logger.debug(lambda: "Evaluated env markers from:\n{}\n\nTo:\n{}".format(
            reqs_with_env_markers,
            env_marker_target_platforms,
        ))

    index_urls = {}
    if get_index_urls:
        index_urls = get_index_urls(
            ctx,
            # Use list({}) as a way to have a set
            list({
                req.distribution: None
                for reqs in requirements_by_platform.values()
                for req in reqs.values()
                if req.srcs.shas
            }),
        )

    ret = {}
    for whl_name, reqs in sorted(requirements_by_platform.items()):
        requirement_target_platforms = {}
        for r in reqs.values():
            target_platforms = env_marker_target_platforms.get(r.requirement_line, r.target_platforms)
            for p in target_platforms:
                requirement_target_platforms[p] = None

        is_exposed = len(requirement_target_platforms) == len(requirements)
        if not is_exposed and logger:
            logger.debug(lambda: "Package '{}' will not be exposed because it is only present on a subset of platforms: {} out of {}".format(
                whl_name,
                sorted(requirement_target_platforms),
                sorted(requirements),
            ))

        # Return normalized names
        ret_requirements = ret.setdefault(normalize_name(whl_name), [])

        for r in sorted(reqs.values(), key = lambda r: r.requirement_line):
            whls, sdist = _add_dists(
                requirement = r,
                index_urls = index_urls.get(whl_name),
                logger = logger,
            )

            target_platforms = env_marker_target_platforms.get(r.requirement_line, r.target_platforms)
            ret_requirements.append(
                struct(
                    distribution = r.distribution,
                    srcs = r.srcs,
                    target_platforms = sorted(target_platforms),
                    extra_pip_args = r.extra_pip_args,
                    whls = whls,
                    sdist = sdist,
                    is_exposed = is_exposed,
                ),
            )

    if logger:
        logger.debug(lambda: "Will configure whl repos: {}".format(ret.keys()))

    return ret

def select_requirement(requirements, *, platform):
    """A simple function to get a requirement for a particular platform.

    Only used in WORKSPACE.

    Args:
        requirements (list[struct]): The list of requirements as returned by
            the `parse_requirements` function above.
        platform (str or None): The host platform. Usually an output of the
            `host_platform` function. If None, then this function will return
            the first requirement it finds.

    Returns:
        None if not found or a struct returned as one of the values in the
        parse_requirements function. The requirement that should be downloaded
        by the host platform will be returned.
    """
    maybe_requirement = [
        req
        for req in requirements
        if not platform or [p for p in req.target_platforms if p.endswith(platform)]
    ]
    if not maybe_requirement:
        # Sometimes the package is not present for host platform if there
        # are whls specified only in particular requirements files, in that
        # case just continue, however, if the download_only flag is set up,
        # then the user can also specify the target platform of the wheel
        # packages they want to download, in that case there will be always
        # a requirement here, so we will not be in this code branch.
        return None

    return maybe_requirement[0]

def host_platform(ctx):
    """Return a string representation of the repository OS.

    Only used in WORKSPACE.

    Args:
        ctx (struct): The `module_ctx` or `repository_ctx` attribute.

    Returns:
        The string representation of the platform that we can later used in the `pip`
        machinery.
    """
    return "{}_{}".format(
        repo_utils.get_platforms_os_name(ctx),
        repo_utils.get_platforms_cpu_name(ctx),
    )

def _add_dists(*, requirement, index_urls, logger = None):
    """Populate dists based on the information from the PyPI index.

    This function will modify the given requirements_by_platform data structure.

    Args:
        requirement: The result of parse_requirements function.
        index_urls: The result of simpleapi_download.
        logger: A logger for printing diagnostic info.
    """

    # Handle direct URLs in requirements
    if requirement.srcs.url:
        url = requirement.srcs.url
        _, _, filename = url.rpartition("/")
        direct_url_dist = struct(
            url = url,
            filename = filename,
            sha256 = requirement.srcs.shas[0] if requirement.srcs.shas else "",
            yanked = False,
        )

        if filename.endswith(".whl"):
            return [direct_url_dist], None
        else:
            return [], direct_url_dist

    if not index_urls:
        return [], None

    whls = []
    sdist = None

    # TODO @aignas 2024-05-22: it is in theory possible to add all
    # requirements by version instead of by sha256. This may be useful
    # for some projects.
    for sha256 in requirement.srcs.shas:
        # For now if the artifact is marked as yanked we just ignore it.
        #
        # See https://packaging.python.org/en/latest/specifications/simple-repository-api/#adding-yank-support-to-the-simple-api

        maybe_whl = index_urls.whls.get(sha256)
        if maybe_whl and not maybe_whl.yanked:
            whls.append(maybe_whl)
            continue

        maybe_sdist = index_urls.sdists.get(sha256)
        if maybe_sdist and not maybe_sdist.yanked:
            sdist = maybe_sdist
            continue

        if logger:
            logger.warn(lambda: "Could not find a whl or an sdist with sha256={}".format(sha256))

    yanked = {}
    for dist in whls + [sdist]:
        if dist and dist.yanked:
            yanked.setdefault(dist.yanked, []).append(dist.filename)
    if yanked:
        logger.warn(lambda: "\n".join([
            "the following distributions got yanked:",
        ] + [
            "reason: {}\n  {}".format(reason, "\n".join(sorted(dists)))
            for reason, dists in yanked.items()
        ]))

    # Filter out the wheels that are incompatible with the target_platforms.
    whls = select_whls(whls = whls, want_platforms = requirement.target_platforms, logger = logger)

    return whls, sdist