Docker image version tag

src/python/pants/backend/docker/docker_build.py

@@ -5,9 +5,15 @@
 import logging
 from dataclasses import dataclass
 from os import path
+from typing import cast
 
 from pants.backend.docker.docker_binary import DockerBinary, DockerBinaryRequest
-from pants.backend.docker.docker_build_context import DockerBuildContext, DockerBuildContextRequest
+from pants.backend.docker.docker_build_context import (
+    DockerBuildContext,
+    DockerBuildContextRequest,
+    DockerVersionContextError,
+    DockerVersionContextValue,
+)
 from pants.backend.docker.registries import DockerRegistries
 from pants.backend.docker.subsystem import DockerOptions
 from pants.backend.docker.target_types import (
@@ -23,11 +29,20 @@
 from pants.engine.process import Process, ProcessResult
 from pants.engine.rules import Get, MultiGet, collect_rules, rule
 from pants.engine.unions import UnionRule
+from pants.util.frozendict import FrozenDict
 from pants.util.strutil import bullet_list, pluralize
 
 logger = logging.getLogger(__name__)
 
 
+class DockerBuildError(Exception):
+    pass
+
+
+class DockerNameTemplateError(ValueError):
+    pass
+
+
 @dataclass(frozen=True)
 class DockerFieldSet(PackageFieldSet):
     required_fields = (DockerImageSources,)
@@ -51,24 +66,51 @@ def dockerfile_path(self) -> str:
         return path.join(self.address.spec_path, self.dockerfile_relpath)
 
     def image_names(
-        self, default_name_template: str, registries: DockerRegistries
+        self,
+        default_name_template: str,
+        registries: DockerRegistries,
+        version_context: FrozenDict[str, DockerVersionContextValue],
     ) -> tuple[str, ...]:
         """This method will always return a non-empty tuple."""
         default_parent = path.basename(path.dirname(self.address.spec_path))
         default_repo = path.basename(self.address.spec_path)
         repo = self.repository.value or default_repo
         name_template = self.name_template.value or default_name_template
-        image_name = name_template.format(
-            name=self.name.value or self.address.target_name,
-            repository=repo,
-            sub_repository="/".join(
-                [default_repo if self.repository.value else default_parent, repo]
-            ),
-        )
-        image_names = tuple(
-            ":".join(s for s in [image_name, tag] if s)
-            for tag in [self.version.value, *(self.tags.value or [])]
-        )
+        try:
+            image_name = name_template.format(
+                name=self.name.value or self.address.target_name,
+                repository=repo,
+                sub_repository="/".join(
+                    [default_repo if self.repository.value else default_parent, repo]
+                ),
+            )
+        except KeyError as e:
+            raise DockerNameTemplateError(
+                f"Invalid image name template: {name_template!r}. Unknown key: {e}.\n\n"
+                f"Use any of 'name', 'repository' or 'sub_repository' in the template string."
+            ) from e
+
+        try:
+            image_names = tuple(
+                ":".join(s for s in [image_name, tag] if s)
+                for tag in [
+                    cast(str, self.version.value).format(**version_context),
+                    *(self.tags.value or []),
+                ]
+            )
+        except (KeyError, ValueError) as e:
+            msg = (
+                "Invalid format string for the `docker_image(version)` field: "
+                f"{self.version.value!r}.\n\n"
+            )
+            if isinstance(e, KeyError):
+                msg += (
+                    f"The key {e} is unknown. Try with one of: "
+                    f'{", ".join(version_context.keys())}.'
+                )
+            else:
+                msg += str(e)
+            raise DockerVersionContextError(msg) from e
 
         registries_options = tuple(registries.get(*(self.registries.value or [])))
         if not registries_options:
@@ -97,7 +139,15 @@ async def build_docker_image(
         ),
     )
 
-    tags = field_set.image_names(options.default_image_name_template, options.registries())
+    try:
+        tags = field_set.image_names(
+            default_name_template=options.default_image_name_template,
+            registries=options.registries(),
+            version_context=context.version_context,
+        )
+    except ValueError as e:
+        raise DockerBuildError(f"Error in {field_set.address}: {e}") from e
+
     result = await Get(
         ProcessResult,
         Process,

src/python/pants/backend/docker/docker_build_context.py

@@ -5,6 +5,7 @@
 from dataclasses import dataclass
 from itertools import chain
 
+from pants.backend.docker.dockerfile_parser import DockerfileInfo
 from pants.backend.docker.target_types import DockerImageSources
 from pants.core.goals.package import BuiltPackage, PackageFieldSet
 from pants.core.target_types import FilesSources
@@ -22,13 +23,31 @@
     TransitiveTargets,
     TransitiveTargetsRequest,
 )
+from pants.util.frozendict import FrozenDict
 
 logger = logging.getLogger(__name__)
 
 
+class DockerVersionContextError(ValueError):
+    pass
+
+
+class DockerVersionContextValue(FrozenDict[str, str]):
+    """Dict class suitable for use as a format string context object, as it allows to use attribute
+    access rather than item access."""
+
+    def __getattr__(self, attribute: str) -> str:
+        if attribute not in self:
+            raise DockerVersionContextError(
+                f"The key {attribute!r} is unknown. Try with one of: " f'{", ".join(self.keys())}.'
+            )
+        return self[attribute]
+
+
 @dataclass(frozen=True)
 class DockerBuildContext:
     digest: Digest
+    version_context: FrozenDict[str, DockerVersionContextValue]
 
 
 @dataclass(frozen=True)
@@ -71,10 +90,11 @@ async def create_docker_build_context(request: DockerBuildContextRequest) -> Doc
         FieldSetsPerTargetRequest(PackageFieldSet, transitive_targets.dependencies),
     )
 
-    dockerfiles, sources, embedded_pkgs_per_target = await MultiGet(
+    dockerfiles, sources, embedded_pkgs_per_target, dockerfile_info = await MultiGet(
         dockerfiles_request,
         sources_request,
         embedded_pkgs_per_target_request,
+        Get(DockerfileInfo, DockerImageSources, transitive_targets.roots[0][DockerImageSources]),
     )
 
     # Package binary dependencies for build context.
@@ -98,7 +118,18 @@ async def create_docker_build_context(request: DockerBuildContextRequest) -> Doc
         MergeDigests(d for d in all_digests if d),
     )
 
-    return DockerBuildContext(context)
+    version_context: dict[str, DockerVersionContextValue] = {}
+    for stage, tag in [tag.split(maxsplit=1) for tag in dockerfile_info.version_tags]:
+        value = DockerVersionContextValue({"tag": tag})
+        if not version_context:
+            # Refer to the first FROM directive as the "baseimage".
+            version_context["baseimage"] = value
+        version_context[stage] = value
+
+    return DockerBuildContext(
+        digest=context,
+        version_context=FrozenDict(version_context),
+    )
 
 
 def rules():

src/python/pants/backend/docker/docker_build_context_test.py

@@ -7,8 +7,13 @@
 
 import pytest
 
-from pants.backend.docker.docker_build_context import DockerBuildContext, DockerBuildContextRequest
+from pants.backend.docker.docker_build_context import (
+    DockerBuildContext,
+    DockerBuildContextRequest,
+    DockerVersionContextValue,
+)
 from pants.backend.docker.docker_build_context import rules as context_rules
+from pants.backend.docker.dockerfile_parser import rules as parser_rules
 from pants.backend.docker.target_types import DockerImage
 from pants.backend.python import target_types_rules
 from pants.backend.python.goals import package_pex_binary
@@ -21,28 +26,33 @@
 from pants.engine.addresses import Address
 from pants.engine.fs import Snapshot
 from pants.testutil.rule_runner import QueryRule, RuleRunner
+from pants.util.frozendict import FrozenDict
 
 
 @pytest.fixture
 def rule_runner() -> RuleRunner:
-    return RuleRunner(
+    rule_runner = RuleRunner(
         rules=[
             *context_rules(),
             *core_target_types_rules(),
             *package_pex_binary.rules(),
+            *parser_rules(),
             *pex_from_targets.rules(),
             *target_types_rules.rules(),
             QueryRule(BuiltPackage, [PexBinaryFieldSet]),
             QueryRule(DockerBuildContext, (DockerBuildContextRequest,)),
         ],
         target_types=[DockerImage, Files, PexBinary],
     )
+    rule_runner.set_options([], env_inherit={"PATH", "PYENV_ROOT", "HOME"})
+    return rule_runner
 
 
 def assert_build_context(
     rule_runner: RuleRunner,
     address: Address,
     expected_files: list[str],
+    expected_version_context: FrozenDict[str, DockerVersionContextValue] | None = None,
 ) -> None:
     context = rule_runner.request(
         DockerBuildContext,
@@ -56,6 +66,8 @@ def assert_build_context(
 
     snapshot = rule_runner.request(Snapshot, [context.digest])
     assert sorted(expected_files) == sorted(snapshot.files)
+    if expected_version_context is not None:
+        assert expected_version_context == context.version_context
 
 
 def test_file_dependencies(rule_runner: RuleRunner) -> None:
@@ -159,7 +171,6 @@ def test_files_out_of_tree(rule_runner: RuleRunner) -> None:
 def test_packaged_pex_path(rule_runner: RuleRunner) -> None:
     # This test is here to ensure that we catch if there is any change in the generated path where
     # built pex binaries go, as we rely on that for dependency inference in the Dockerfile.
-    rule_runner.set_options([], env_inherit={"PATH", "PYENV_ROOT", "HOME"})
     rule_runner.write_files(
         {
             "src/docker/BUILD": """docker_image(dependencies=["src/python/proj/cli:bin"])""",
@@ -174,3 +185,34 @@ def test_packaged_pex_path(rule_runner: RuleRunner) -> None:
         Address("src/docker", target_name="docker"),
         expected_files=["src/docker/Dockerfile", "src.python.proj.cli/bin.pex"],
     )
+
+
+def test_version_context_from_dockerfile(rule_runner: RuleRunner) -> None:
+    rule_runner.write_files(
+        {
+            "src/docker/BUILD": """docker_image()""",
+            "src/docker/Dockerfile": dedent(
+                """\
+                FROM python:3.8
+                FROM alpine as interim
+                FROM interim
+                FROM scratch:1-1 as output
+                """
+            ),
+        }
+    )
+
+    assert_build_context(
+        rule_runner,
+        Address("src/docker"),
+        expected_files=["src/docker/Dockerfile"],
+        expected_version_context=FrozenDict(
+            {
+                "baseimage": DockerVersionContextValue({"tag": "3.8"}),
+                "stage0": DockerVersionContextValue({"tag": "3.8"}),
+                "interim": DockerVersionContextValue({"tag": "latest"}),
+                "stage2": DockerVersionContextValue({"tag": "latest"}),
+                "output": DockerVersionContextValue({"tag": "1-1"}),
+            }
+        ),
+    )