"""The transition module contains the rule definitions to wrap py_binary and py_test and transition

them to the desired target platform.

"""

load("@bazel_skylib//lib:dicts.bzl", "dicts")

load("//python:defs.bzl", _py_binary = "py_binary", _py_test = "py_test")

def _transition_python_version_impl(_, attr):

return {"//python/config_settings:python_version": str(attr.python_version)}

_transition_python_version = transition(

implementation = _transition_python_version_impl,

inputs = [],

outputs = ["//python/config_settings:python_version"],

)

def _transition_py_impl(ctx):

target = ctx.attr.target

windows_constraint = ctx.attr._windows_constraint[platform_common.ConstraintValueInfo]

target_is_windows = ctx.target_platform_has_constraint(windows_constraint)

executable = ctx.actions.declare_file(ctx.attr.name + (".exe" if target_is_windows else ""))

ctx.actions.symlink(

is_executable = True,

output = executable,

target_file = target[DefaultInfo].files_to_run.executable,

)

zipfile_symlink = None

if target_is_windows:

# Under Windows, the expected "<name>.zip" does not exist, so we have to

# create the symlink ourselves to achieve the same behaviour as in macOS

# and Linux.

zipfile = None

expected_target_path = target[DefaultInfo].files_to_run.executable.short_path[:-4] + ".zip"

for file in target[DefaultInfo].default_runfiles.files.to_list():

if file.short_path == expected_target_path:

zipfile = file

zipfile_symlink = ctx.actions.declare_file(ctx.attr.name + ".zip")

ctx.actions.symlink(

is_executable = True,

output = zipfile_symlink,

target_file = zipfile,

)

env = {}

for k, v in ctx.attr.env.items():

env[k] = ctx.expand_location(v)

providers = [

DefaultInfo(

executable = executable,

files = depset([zipfile_symlink] if zipfile_symlink else [], transitive = [target[DefaultInfo].files]),

runfiles = ctx.runfiles([zipfile_symlink] if zipfile_symlink else []).merge(target[DefaultInfo].default_runfiles),

),

target[PyInfo],

target[PyRuntimeInfo],

# Ensure that the binary we're wrapping is included in code coverage.

coverage_common.instrumented_files_info(

ctx,

dependency_attributes = ["target"],

),

target[OutputGroupInfo],

# TODO(f0rmiga): testing.TestEnvironment is deprecated in favour of RunEnvironmentInfo but

# RunEnvironmentInfo is not exposed in Bazel < 5.3.

# https://github.com/bazelbuild/rules_python/issues/901

# https://github.com/bazelbuild/bazel/commit/dbdfa07e92f99497be9c14265611ad2920161483

testing.TestEnvironment(env),

]

return providers

_COMMON_ATTRS = {

"deps": attr.label_list(

mandatory = False,

),

"env": attr.string_dict(

mandatory = False,

),

"python_version": attr.string(

mandatory = True,

),

"srcs": attr.label_list(

allow_files = True,

mandatory = False,

),

"target": attr.label(

executable = True,

cfg = "target",

mandatory = True,

providers = [PyInfo],

),

# "tools" is a hack here. It should be "data" but "data" is not included by default in the

# location expansion in the same way it is in the native Python rules. The difference on how

# the Bazel deals with those special attributes differ on the LocationExpander, e.g.:

# https://github.com/bazelbuild/bazel/blob/ce611646/src/main/java/com/google/devtools/build/lib/analysis/LocationExpander.java#L415-L429

#

# Since the default LocationExpander used by ctx.expand_location is not the same as the native

# rules (it doesn't set "allowDataAttributeEntriesInLabel"), we use "tools" temporarily while a

# proper fix in Bazel happens.

#

# A fix for this was proposed in https://github.com/bazelbuild/bazel/pull/16381.

"tools": attr.label_list(

allow_files = True,

mandatory = False,

),

# Required to Opt-in to the transitions feature.

"_allowlist_function_transition": attr.label(

default = "@bazel_tools//tools/allowlists/function_transition_allowlist",

),

"_windows_constraint": attr.label(

default = "@platforms//os:windows",

),

}

_transition_py_binary = rule(

_transition_py_impl,

attrs = _COMMON_ATTRS,

cfg = _transition_python_version,

executable = True,

)

_transition_py_test = rule(

_transition_py_impl,

attrs = _COMMON_ATTRS,

cfg = _transition_python_version,

test = True,

)

def _py_rule(rule_impl, transition_rule, name, python_version, **kwargs):

args = kwargs.pop("args", None)

data = kwargs.pop("data", None)

env = kwargs.pop("env", None)

srcs = kwargs.pop("srcs", None)

deps = kwargs.pop("deps", None)

# Attributes common to all build rules.

# https://bazel.build/reference/be/common-definitions#common-attributes

compatible_with = kwargs.pop("compatible_with", None)

deprecation = kwargs.pop("deprecation", None)

distribs = kwargs.pop("distribs", None)

exec_compatible_with = kwargs.pop("exec_compatible_with", None)

exec_properties = kwargs.pop("exec_properties", None)

features = kwargs.pop("features", None)

restricted_to = kwargs.pop("restricted_to", None)

tags = kwargs.pop("tags", None)

target_compatible_with = kwargs.pop("target_compatible_with", None)

testonly = kwargs.pop("testonly", None)

toolchains = kwargs.pop("toolchains", None)

visibility = kwargs.pop("visibility", None)

common_attrs = {

"compatible_with": compatible_with,

"deprecation": deprecation,

"distribs": distribs,

"exec_compatible_with": exec_compatible_with,

"exec_properties": exec_properties,

"features": features,

"restricted_to": restricted_to,

"target_compatible_with": target_compatible_with,

"testonly": testonly,

"toolchains": toolchains,

}

# Test-specific extra attributes.

if "env_inherit" in kwargs:

common_attrs["env_inherit"] = kwargs.pop("env_inherit")

if "size" in kwargs:

common_attrs["size"] = kwargs.pop("size")

if "timeout" in kwargs:

common_attrs["timeout"] = kwargs.pop("timeout")

if "flaky" in kwargs:

common_attrs["flaky"] = kwargs.pop("flaky")

if "shard_count" in kwargs:

common_attrs["shard_count"] = kwargs.pop("shard_count")

if "local" in kwargs:

common_attrs["local"] = kwargs.pop("local")

# Binary-specific extra attributes.

if "output_licenses" in kwargs:

common_attrs["output_licenses"] = kwargs.pop("output_licenses")

rule_impl(

name = "_" + name,

args = args,

data = data,

deps = deps,

env = env,

srcs = srcs,

tags = ["manual"] + (tags if tags else []),

visibility = ["//visibility:private"],

**dicts.add(common_attrs, kwargs)

)

return transition_rule(

name = name,

args = args,

deps = deps,

env = env,

python_version = python_version,

srcs = srcs,

tags = tags,

target = ":_" + name,

tools = data,

visibility = visibility,

**common_attrs

)

def py_binary(name, python_version, **kwargs):

return _py_rule(_py_binary, _transition_py_binary, name, python_version, **kwargs)

def py_test(name, python_version, **kwargs):

return _py_rule(_py_test, _transition_py_test, name, python_version, **kwargs)