plugin.py

from dataclasses import dataclass
import logging
from pathlib import Path

from pants.core.goals.package import PackageFieldSet, BuiltPackage
from pants.core.util_rules.source_files import SourceFiles, SourceFilesRequest
from pants.core.util_rules.stripped_source_files import StrippedSourceFiles
from pants.engine.addresses import Address, AddressInput
from pants.engine.fs import EMPTY_DIGEST, MergeDigests, Snapshot
from pants.engine.process import (
    BinaryPathRequest,
    BinaryPathTest,
    BinaryPaths,
    Process,
    ProcessResult,
)
from pants.engine.rules import collect_rules, Get, MultiGet, rule
from pants.engine.target import (
    COMMON_TARGET_FIELDS,
    Dependencies,
    FieldSetsPerTarget,
    FieldSetsPerTargetRequest,
    Sources,
    Target,
    Targets,
    StringField,
    StringSequenceField,
    TransitiveTargets,
    WrappedTarget,
)
from pants.engine.unions import UnionRule
from pants.util.logging import LogLevel


logger = logging.getLogger(__name__)


class DockerfileSources(Sources):
    default = ['Dockerfile']
    expected_num_files = 1


class DockerImageRepoName(StringField):
    alias = "image_repo"
    required = True


class DockerImageTagName(StringField):
    alias = "image_tag"
    default = "latest"


class DockerPackages(StringSequenceField):
    alias = "packages"


class DockerImage(Target):
    alias = "docker_image"
    core_fields = (
        *COMMON_TARGET_FIELDS,
        Dependencies,
        DockerfileSources,
        DockerImageRepoName,
        DockerImageTagName,
        DockerPackages,
    )


@dataclass(frozen=True)
class DockerImageFieldSet(PackageFieldSet):
    required_fields = (DockerfileSources, DockerImageRepoName)

    sources: DockerfileSources
    image_repo: DockerImageRepoName
    packages: DockerPackages


@rule(level=LogLevel.DEBUG)
async def create_docker_image(field_set: DockerImageFieldSet) -> BuiltPackage:
    docker_paths = await Get(
        BinaryPaths,
        BinaryPathRequest(
            binary_name="docker",
            # TODO: why not just $PATH variable?
            search_path=["/usr/local/bin", "/usr/bin", "/bin"],
            test=BinaryPathTest(args=["--version"]),
        ),
    )
    docker_bin = docker_paths.first_path
    if docker_bin is None:
        raise OSError("could not find 'docker'")

    dockerfile_sources = await Get(SourceFiles, SourceFilesRequest([field_set.sources]))

    # TODO: swap this in when UnparsedAddressInputs becomes available
    # package_targets = await Get(
    #     Targets,
    #     UnparsedAddressInputs(field_set.packages.value or (), owning_address=field_set.address),
    # )
    wrapped_package_targets = await MultiGet(
        Get(
            WrappedTarget,
            AddressInput,
            AddressInput.parse(spec, relative_to=field_set.address.spec_path),
        )
        for spec in field_set.packages.value or ()
    )
    package_targets = (wt.target for wt in wrapped_package_targets)
    package_field_sets_per_target = await Get(
        FieldSetsPerTarget, FieldSetsPerTargetRequest(PackageFieldSet, package_targets)
    )
    packages = await MultiGet(
        Get(BuiltPackage, PackageFieldSet, field_set)
        for field_set in package_field_sets_per_target.field_sets
    )

    input_snapshot = await Get(
        Snapshot,
        MergeDigests((
            dockerfile_sources.snapshot.digest,
            *(p.digest for p in packages),
        )),
    )

    options = [
        '-f', dockerfile_sources.snapshot.files[0],
        '-t', field_set.image_repo.value,
    ]

    result = await Get(
        ProcessResult,
        Process(
            argv=[docker_bin.path, 'build', *options, '.'],
            input_digest=input_snapshot.digest,
            description=f"build {field_set.address} into a docker image",
        ),
    )
    return BuiltPackage(EMPTY_DIGEST, relpath="")



def rules():
    return [
        *collect_rules(),
        UnionRule(PackageFieldSet, DockerImageFieldSet),
    ]