set_docker_image_tags.py

from pants.backend.docker.target_types import DockerImageTags, DockerImageTagsRequest
from pants.engine.env_vars import EnvironmentVars, EnvironmentVarsRequest
from pants.engine.internals.selectors import Get
from pants.engine.rules import collect_rules, rule
from pants.engine.target import Target
from pants.engine.unions import UnionRule
from pants_utils.get_version_tag import get_version_tag


class CustomDockerImageTagsRequest(DockerImageTagsRequest):
    @classmethod
    def is_applicable(cls, target: Target) -> bool:
        """
        Apply this custom rule to all docker_image targets.
        """
        return True


@rule
async def custom_image_tags(request: CustomDockerImageTagsRequest) -> DockerImageTags:
    """
    Adds additional tags to Docker images.
    """

    # Check if the USE_VERSION_IMAGE_TAGS env var is set
    env = await Get(EnvironmentVars, EnvironmentVarsRequest(["USE_VERSION_IMAGE_TAGS"]))
    use_version_img_tags = str(env.get("USE_VERSION_IMAGE_TAGS"))

    # Return empty tags if the env var is not set
    if use_version_img_tags is None or use_version_img_tags.lower() != "true":
        return DockerImageTags([])
    version = get_version_tag(request.target.address.spec_path)
    return DockerImageTags([version] if version is not None else [])


def rules():
    return (
        *collect_rules(),
        UnionRule(DockerImageTagsRequest, CustomDockerImageTagsRequest),
    )

register.py

from pants.engine.rules import collect_rules
from set_docker_image_tags.set_docker_image_tags import rules as docker_image_tag_rules


def rules():
    return (
        *collect_rules(),
        *docker_image_tag_rules(),
    )

test_set_docker_image_tags.py

from textwrap import dedent

from pants.backend.docker.goals.package_image import rules as package_image_rules
from pants.backend.docker.subsystems.docker_options import DockerOptions
from pants.backend.docker.target_types import DockerImageTags, DockerImageTarget
from pants.backend.docker.util_rules.docker_build_args import (
    DockerBuildArgs,
    DockerBuildArgsRequest,
)
from pants.backend.docker.util_rules.docker_build_args import rules as build_args_rules
from pants.backend.docker.util_rules.docker_build_env import (
    DockerBuildEnvironment,
    DockerBuildEnvironmentRequest,
)
from pants.backend.docker.util_rules.docker_build_env import rules as build_env_rules
from pants.engine.addresses import Address
from pants.engine.env_vars import EnvironmentVars, EnvironmentVarsRequest
from pants.option.global_options import GlobalOptions
from pants.testutil.rule_runner import QueryRule, RuleRunner
from pants.util.frozendict import FrozenDict
from set_docker_image_tags.set_docker_image_tags import CustomDockerImageTagsRequest


def test_example() -> None:
    rule_runner = RuleRunner(
        target_types=[DockerImageTarget],
        rules=[
            *package_image_rules(),
            *build_args_rules(),
            *build_env_rules(),
            QueryRule(GlobalOptions, []),
            QueryRule(DockerOptions, []),
            QueryRule(DockerBuildArgs, [DockerBuildArgsRequest]),
            QueryRule(DockerBuildEnvironment, [DockerBuildEnvironmentRequest]),
            QueryRule(DockerImageTags, [CustomDockerImageTagsRequest]),
        ],
    )
    rule_runner.write_files(
        {
            "project/BUILD": dedent(
                """\
             docker_image()
             """
            )
        }
    )
    tgt = rule_runner.get_target(Address("project"))

    def mock_env_vars(request: EnvironmentVarsRequest) -> EnvironmentVars:
        return EnvironmentVars(FrozenDict({"USE_VERSION_IMAGE_TAGS": "true"}))

    custom_docker_img_tags_request = CustomDockerImageTagsRequest(target=tgt)
    result = rule_runner.request(DockerImageTags, [custom_docker_img_tags_request])

    # TODO: add proper assert stmt
    print(result)
    assert True

E.     ValueError: Encountered 1 rule graph error:
E         No installed rules return the type DockerImageTags, and it was not provided by potential callers of Query(DockerImageTags for (CustomDockerImageTagsRequest, EnvironmentName)).
E           If that type should be computed by a rule, ensure that that rule is installed.
E           If it should be provided by a caller, ensure that it is included in any relevant Query or Get.

custom_image_tags

set_docker_image_tags.rules

run_pants

bootstrap_args

rule_runner.run_goal_rule