register.py

"""Version tag plugin."""

import logging
import os
from typing import Any, cast

from pants.backend.python.target_types import PythonSourceField
from pants.engine.fs import (
    AddPrefix,
    CreateDigest,
    Digest,
    FileContent,
    Snapshot,
)
from pants.engine.rules import Get, collect_rules, rule
from pants.engine.target import (
    COMMON_TARGET_FIELDS,
    Dependencies,
    GeneratedSources,
    GenerateSourcesRequest,
    OptionalSingleSourceField,
    StringField,
    Target,
)
from pants.engine.unions import UnionRule
from pants.source.source_root import SourceRoot, SourceRootRequest
from pants.util.logging import LogLevel

logger = logging.getLogger(__name__)


class VersionFileSourceField(OptionalSingleSourceField):
    """Version file to write to."""

    default = "version.py"
    help = "File to store version (eg. `version.py`)."
    # required = True


class VersionField(StringField):
    """Version to write to file."""

    alias = "version_tag"
    help = "Version to write to file."
    required = True


class VersionFileSourceTarget(Target):
    """A single version file target."""

    alias = "write_version_file"
    help = "A single version file."
    core_fields = (*COMMON_TARGET_FIELDS, Dependencies, VersionFileSourceField)


class GeneratePythonVersionFileRequest(GenerateSourcesRequest):
    """Generate a Python version file from a version file target."""

    input = VersionFileSourceField
    output = PythonSourceField


@rule(level=LogLevel.DEBUG)
async def write_version_file_func(
    request: GeneratePythonVersionFileRequest,
) -> GeneratedSources:
    """Write a version file from a version file target.

    Parameters
    ----------
    request
        The request.
    """
    package_name = request.protocol_target.address.spec_path
    output_filename = cast(str, request.protocol_target[VersionFileSourceField].value)
    version_tag = request.protocol_target[VersionField].value

    logger.debug(f"package_name: {package_name}")
    logger.debug(f"output_filename: {output_filename}")
    logger.debug(f"version_tag: {version_tag}")

    digest = await Get(
        Digest,
        CreateDigest(
            [
                FileContent(
                    os.path.join(package_name, output_filename),
                    f"version = '{version_tag}'\n".encode(),
                )
            ]
        ),
    )

    source_root = await Get(
        SourceRoot,
        SourceRootRequest,
        SourceRootRequest.for_target(request.protocol_target),
    )

    source_root_restored = (
        await Get(Snapshot, AddPrefix(digest, source_root.path))
        if source_root.path != "."
        else await Get(Snapshot, Digest, digest)
    )
    return GeneratedSources(source_root_restored)


def target_types() -> list[type[Target]]:
    """Return the list of target types this plugin registers."""
    return [VersionFileSourceTarget]


def rules() -> Any:
    """Return the list of rules this plugin registers."""
    return [
        *collect_rules(),
        UnionRule(GenerateSourcesRequest, GeneratePythonVersionFileRequest),
        VersionFileSourceTarget.register_plugin_field(VersionField),
    ]

BUILD

write_version_file(
    name="version",
    source="version.py",
    version_tag=env("BUILD_TAG", "undefined_build_tag_env"),
)

python_sources(
    name="test_package",
    dependencies=[":version"],
)

GenerateSourcesRequest

VCSVersionDummySourceField

OptionalSourceField

source

alias

_