fix typechecks

src/python/pants/backend/docker/lint/trivy/rules.py

@@ -1,8 +1,9 @@
 # Copyright 2024 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, cast
 
+from pants.backend.docker.package_types import BuiltDockerImage
 from pants.backend.docker.target_types import DockerImageSourceField, DockerImageTarget
 from pants.backend.tools.trivy.rules import RunTrivyRequest, run_trivy
 from pants.backend.tools.trivy.subsystem import SkipTrivyField, Trivy
@@ -12,7 +13,7 @@
 from pants.engine.addresses import Addresses
 from pants.engine.internals.native_engine import EMPTY_DIGEST
 from pants.engine.internals.selectors import Get
-from pants.engine.rules import collect_rules, rule
+from pants.engine.rules import collect_rules, rule, implicitly
 from pants.engine.target import (
     FieldSet,
     FieldSetsPerTarget,
@@ -52,11 +53,9 @@ def command_args():
 
 @rule(desc="Lint Docker image with Trivy", level=LogLevel.DEBUG)
 async def run_trivy_docker(
-    request: TrivyDockerRequest.Batch[TrivyDockerRequest, Any],
+    request: TrivyDockerRequest.Batch[TrivyDockerFieldSet, Any],
 ) -> LintResult:
-    assert len(request.elements) == 1, "not single element in partition"  # "Do we need to?"
     addrs = tuple(e.address for e in request.elements)
-
     tgts = await Get(Targets, Addresses(addrs))
 
     field_sets_per_tgt = await Get(
@@ -65,15 +64,17 @@ async def run_trivy_docker(
     [field_set] = field_sets_per_tgt.field_sets
 
     package = await Get(BuiltPackage, EnvironmentAwarePackageRequest(field_set))
+    built_image: BuiltDockerImage = cast(BuiltDockerImage, package.artifacts[0])
     r = await run_trivy(
         RunTrivyRequest(
             command="image",
             command_args=command_args(),
             scanners=(),
-            target=package.artifacts[0].image_id,
+            target=built_image.image_id,
             input_digest=EMPTY_DIGEST,
-            description=f"Run Trivy on docker image {','.join(package.artifacts[0].tags)}",
-        )
+            description=f"Run Trivy on docker image {','.join(built_image.tags)}",
+        ),
+        **implicitly(),
     )
 
     return LintResult.create(request, r)

src/python/pants/backend/helm/lint/trivy/rules.py

@@ -24,7 +24,7 @@
 from pants.core.util_rules.partitions import PartitionerType
 from pants.engine.internals.selectors import Get
 from pants.engine.process import FallibleProcessResult
-from pants.engine.rules import collect_rules, rule
+from pants.engine.rules import collect_rules, rule, implicitly
 from pants.engine.target import FieldSet, Target
 from pants.util.logging import LogLevel
 
@@ -61,7 +61,8 @@ async def run_trivy_on_helm(
             target=".",  # the charts are rendered to the local directory
             input_digest=request.rendered_files.snapshot.digest,
             description=f"Run Trivy on Helm files for {request.field_set.address}",
-        )
+        ),
+        **implicitly(),
     )
 
     return r
@@ -114,12 +115,12 @@ class TrivyLintHelmChartRequest(TrivyLintHelmRequest):
 
 @rule(desc="Lint Helm chart with Trivy", level=LogLevel.DEBUG)
 async def run_trivy_on_helm_chart(
-    request: TrivyLintHelmChartRequest.Batch[TrivyLintHelmChartRequest, Any],
+    request: TrivyLintHelmChartRequest.Batch[TrivyLintHelmChartFieldSet, Any],
 ) -> LintResult:
     assert len(request.elements) == 1, "not single element in partition"  # "Do we need to?"
     [field_set] = request.elements
 
-    rendered_files = await Get(RenderedHelmFiles, RenderHelmChartRequest(field_set))
+    rendered_files: RenderedHelmFiles = await Get(RenderedHelmFiles, RenderHelmChartRequest(field_set))
     r = await run_trivy_on_helm(RunTrivyOnHelmRequest(field_set, rendered_files))
 
     return LintResult.create(request, r)

src/python/pants/backend/terraform/lint/trivy/rules.py

@@ -2,7 +2,7 @@
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 from abc import ABCMeta
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Union
 
 from pants.backend.terraform.dependencies import terraform_fieldset_to_init_request, terraform_init
 from pants.backend.terraform.dependency_inference import (
@@ -23,7 +23,7 @@
 from pants.engine.internals.native_engine import MergeDigests
 from pants.engine.intrinsics import merge_digests
 from pants.engine.process import FallibleProcessResult
-from pants.engine.rules import collect_rules, rule
+from pants.engine.rules import collect_rules, rule, implicitly
 from pants.engine.target import FieldSet, SourcesField, Target
 from pants.util.logging import LogLevel
 
@@ -47,7 +47,8 @@ class RunTrivyOnTerraformRequest:
 @rule
 async def run_trivy_on_terraform(req: RunTrivyOnTerraformRequest) -> FallibleProcessResult:
     fs = req.field_set
-    tf = await terraform_init(terraform_fieldset_to_init_request(fs))
+    # Each subclass of TrivyTerraformFieldSet is a subclass of either TerraformDeploymentFieldSet or TerraformFieldSet
+    tf = await terraform_init(terraform_fieldset_to_init_request(fs))  # type: ignore
     command_args = []
 
     if isinstance(fs, TerraformDeploymentFieldSet):
@@ -76,7 +77,8 @@ async def run_trivy_on_terraform(req: RunTrivyOnTerraformRequest) -> FalliblePro
             target=tf.chdir,
             input_digest=input_digest,
             description=f"Run Trivy on terraform deployment {fs.address}",
-        )
+        ),
+        **implicitly(),
     )
 
 
@@ -93,7 +95,7 @@ class TrivyLintTerraformDeploymentRequest(TrivyLintTerraformRequest):
 
 @rule(desc="Lint Terraform deployment with Trivy", level=LogLevel.DEBUG)
 async def run_trivy_on_terraform_deployment(
-    request: TrivyLintTerraformDeploymentRequest.Batch[TrivyLintTerraformDeploymentRequest, Any]
+    request: TrivyLintTerraformDeploymentRequest.Batch[TrivyLintTerraformDeploymentFieldSet, Any]
 ) -> LintResult:
     assert len(request.elements) == 1, "not single element in partition"  # "Do we need to?"
     [fs] = request.elements
@@ -114,7 +116,7 @@ class TrivyLintTerraformModuleRequest(TrivyLintTerraformRequest):
 
 @rule(desc="Lint Terraform module with Trivy", level=LogLevel.DEBUG)
 async def run_trivy_on_terraform_module(
-    request: TrivyLintTerraformModuleRequest.Batch[TrivyLintTerraformModuleRequest, Any]
+    request: TrivyLintTerraformModuleRequest.Batch[TrivyLintTerraformModuleFieldSet, Any]
 ) -> LintResult:
     assert len(request.elements) == 1, "not single element in partition"  # "Do we need to?"
     [fs] = request.elements

src/python/pants/backend/terraform/lint/trivy/trivy_integration_test.py

@@ -20,7 +20,7 @@
 from pants.backend.tools.trivy.testutil import assert_trivy_output, trivy_config
 from pants.core.goals.lint import LintResult
 from pants.core.util_rules import source_files
-from pants.core.util_rules.partitions import PartitionMetadata
+from pants.core.util_rules.partitions import PartitionMetadata, _EmptyMetadata
 from pants.engine.internals.native_engine import Address
 from pants.engine.rules import QueryRule
 from pants.testutil.rule_runner import RuleRunner
@@ -83,7 +83,7 @@ def test_lint_deployment(rule_runner) -> None:
         LintResult,
         [
             TrivyLintTerraformDeploymentRequest.Batch(
-                "trivy", (TerraformDeploymentFieldSet.create(tgt),), PartitionMetadata
+                "trivy", (TerraformDeploymentFieldSet.create(tgt),), partition_metadata=_EmptyMetadata()
             )
         ],
     )
@@ -100,7 +100,7 @@ def test_lint_module(rule_runner) -> None:
         LintResult,
         [
             TrivyLintTerraformModuleRequest.Batch(
-                "trivy", (TerraformFieldSet.create(tgt),), PartitionMetadata
+                "trivy", (TerraformFieldSet.create(tgt),), partition_metadata=_EmptyMetadata()
             )
         ],
     )

src/python/pants/backend/tools/trivy/rules.py

@@ -15,7 +15,7 @@
 from pants.engine.intrinsics import execute_process, merge_digests
 from pants.engine.platform import Platform
 from pants.engine.process import FallibleProcessResult, Process
-from pants.engine.rules import collect_rules, rule
+from pants.engine.rules import collect_rules, rule, implicitly
 from pants.engine.unions import UnionRule
 from pants.option.global_options import KeepSandboxes
 from pants.util.logging import LogLevel
@@ -87,7 +87,8 @@ async def run_trivy(
             env=env,
             description=request.description,
             level=LogLevel.DEBUG,
-        )
+        ),
+        **implicitly(),
     )
     return result