transitions and rename py_binary

py/private/BUILD.bazel

@@ -27,6 +27,7 @@ bzl_library(
     srcs = ["py_image_layer.bzl"],
     deps = [
         "@aspect_bazel_lib//lib:tar",
+        "@aspect_bazel_lib//lib:transitions",
     ],
 )
 

py/private/py_image_layer.bzl

@@ -31,6 +31,7 @@ oci_image(
 """
 
 load("@aspect_bazel_lib//lib:tar.bzl", "mtree_spec", "tar")
+load("@aspect_bazel_lib//lib:transitions.bzl", "platform_transition_filegroup")
 
 default_layer_groups = {
     # match *only* external pip like repositories that contain the string "site-packages"
@@ -90,7 +91,7 @@ awk < $< 'BEGIN {
     )
 
 
-def py_image_layer(name, py_binary, root = "/", layer_groups = {}, compress = "gzip", tar_args = ["--options", "gzip:!timestamp"], compute_unused_inputs = 1, **kwargs):
+def py_image_layer(name, binary, root = "/", layer_groups = {}, compress = "gzip", tar_args = ["--options", "gzip:!timestamp"], compute_unused_inputs = 1, platform = None, **kwargs):
     """Produce a separate tar output for each layer of a python app
 
     > Requires `awk` to be installed on the host machine/rbe runner.
@@ -113,11 +114,12 @@ def py_image_layer(name, py_binary, root = "/", layer_groups = {}, compress = "g
 
     Args:
         name: base name for targets
-        py_binary: a py_binary target
+        binary: a py_binary target
         root: Path to where the layers should be rooted. If not specified, the layers will be rooted at the workspace root.
         layer_groups: Additional layer groups to create. They are used to group files into layers based on their path. In the form of: ```{"<name>": "regex_to_match_against_file_paths"}```
         compress: Compression algorithm to use. Default is gzip. See: https://github.com/bazel-contrib/bazel-lib/blob/main/docs/tar.md#tar_rule-compress
         compute_unused_inputs: Whether to compute unused inputs. Default is 1. See: https://github.com/bazel-contrib/bazel-lib/blob/main/docs/tar.md#tar_rule-compute_unused_inputs
+        platform: The platform to use for the transition. Default is None. See: https://github.com/bazel-contrib/bazel-lib/blob/main/docs/transitions.md#platform_transition_binary-target_platform
         tar_args: Additional arguments to pass to the tar rule. Default is `["--options", "gzip:!timestamp"]`. See: https://github.com/bazel-contrib/bazel-lib/blob/main/docs/tar.md#tar_rule-args
         **kwargs: attribute that apply to all targets expanded by the macro
 
@@ -131,7 +133,7 @@ def py_image_layer(name, py_binary, root = "/", layer_groups = {}, compress = "g
     # into fine-grained layers for better pull, push and remote cache performance.
     mtree_spec(
         name = name + ".manifest",
-        srcs = [py_binary],
+        srcs = [binary],
         **kwargs
     )
 
@@ -147,19 +149,27 @@ def py_image_layer(name, py_binary, root = "/", layer_groups = {}, compress = "g
         tar_target = "_{}_{}".format(name, group_name)
         tar(
             name = tar_target,
-            srcs = [py_binary],
+            srcs = [binary],
             mtree = "{}.{}.manifest.spec".format(name, group_name),
             compress = compress,
             compute_unused_inputs = compute_unused_inputs,
             args = tar_args,
             **kwargs
         )
         srcs.append(tar_target)
-
-    native.filegroup(
-        name = name,
-        srcs = srcs,
-        **kwargs
-    )
+
+    if platform:
+        platform_transition_filegroup(
+            name = name,
+            srcs = srcs,
+            target_platform = platform,
+            **kwargs
+        )
+    else:
+        native.filegroup(
+            name = name,
+            srcs = srcs,
+            **kwargs
+        )
 
     return srcs

py/tests/py_image_layer/BUILD.bazel

@@ -1,6 +1,15 @@
 load("//py:defs.bzl", "py_image_layer", "py_binary")
 load("asserts.bzl", "assert_tar_listing")
 
+platform(
+    name = "linux_amd64",
+    constraint_values = [
+        "@platforms//os:linux",
+        "@platforms//cpu:x86_64",
+    ],
+)
+
+
 # Case 1: Basic usage
 py_binary(
     name = "my_app_bin",
@@ -9,7 +18,8 @@ py_binary(
 
 py_image_layer(
     name = "my_app_layers",
-    py_binary = ":my_app_bin",
+    binary = ":my_app_bin",
+    platform = ":linux_amd64",
 )
 
 assert_tar_listing(