feat: support python_version attribute (#347)

This PR adds support python_version attribute the same way rules_python
does, with the exception that we generate no `defs.bzl` file that users
can just replace their with. Instead we encourage them to create their
own macros that sets the desired version on py_binary/py_test rule.

 
### Changes are visible to end-users: yes

- Searched for relevant documentation and updated as needed: yes
- Breaking change (forces users to change their own code or config): no
- Suggested release notes appear below: yes

### Test plan

- New test cases added

---------

Co-authored-by: Matt Mackay <[email protected]>

MODULE.bazel

@@ -13,6 +13,14 @@ bazel_dep(name = "bazel_skylib", version = "1.4.2")
 bazel_dep(name = "rules_python", version = "0.29.0")
 bazel_dep(name = "platforms", version = "0.0.7")
 
+
+# Custom python version for testing only
+python = use_extension("@rules_python//python/extensions:python.bzl", "python")
+python.toolchain(
+    is_default = False,
+    python_version = "3.8.12",
+)
+
 tools = use_extension("//py:extensions.bzl", "py_tools")
 tools.rules_py_tools()
 use_repo(tools, "rules_py_tools")

WORKSPACE

@@ -24,6 +24,18 @@ register_toolchains("//:container_py_toolchain")
 
 load("@rules_python//python:repositories.bzl", "py_repositories", "python_register_toolchains")
 
+python_register_toolchains(
+    name = "python_toolchain_3_8",
+    python_version = "3.8.12",
+    # Setting `set_python_version_constraint` will set special constraints on the registered toolchain.
+    # This means that this toolchain registration will only be selected for `py_binary` / `py_test` targets 
+    # that have the `python_version = "3.8.12"` attribute set. Targets that have no `python_attribute` will use
+    # the default toolchain resolved which can be seen below.
+    set_python_version_constraint = True,
+)
+
+# It is important to register the default toolchain at last as it will be selected for any
+# py_test/py_binary target even if it has python_version attribute set.
 python_register_toolchains(
     name = "python_toolchain",
     python_version = "3.9",

examples/multi_version/BUILD.bazel

@@ -0,0 +1,38 @@
+load("@aspect_rules_py//py:defs.bzl", "py_binary", "py_pytest_main", "py_test")
+
+py_binary(
+    name = "multi_version",
+    srcs = ["say.py"],
+    deps = [
+        "@pypi_cowsay//:pkg",
+    ],
+    python_version = "3.8.12"
+)
+py_pytest_main(
+    name = "__test__",
+    deps = ["@pypi_pytest//:pkg"],
+)
+
+py_test(
+    name = "py_version_test",
+    srcs = [
+        "py_version_test.py",
+        ":__test__",
+    ],
+    main = ":__test__.py",
+    deps = [
+        ":__test__",
+    ],
+    python_version = "3.8.12"
+)
+py_test(
+    name = "py_version_default_test",
+    srcs = [
+        "py_version_default_test.py",
+        ":__test__",
+    ],
+    main = ":__test__.py",
+    deps = [
+        ":__test__",
+    ],
+)

examples/multi_version/py_version_default_test.py

@@ -0,0 +1,5 @@
+import sys
+
+def test_default_py_version():
+    assert sys.version_info.major == 3, "sys.version_info.major == 3"
+    assert sys.version_info.minor == 9, "sys.version_info.minor == 9"

examples/multi_version/py_version_test.py

@@ -0,0 +1,6 @@
+import sys
+
+def test_specific_py_version():
+    assert sys.version_info.major == 3, "sys.version_info.major == 3"
+    assert sys.version_info.minor == 8, "sys.version_info.minor == 8"
+    assert sys.version_info.micro == 12, "sys.version_info.micro == 12"

examples/multi_version/say.py

@@ -0,0 +1,4 @@
+import cowsay
+import sys
+
+cowsay.cow('hello py_binary, %s!' % sys.version)

py/private/py_binary.bzl

@@ -139,6 +139,41 @@ _attrs = dict({
         allow_single_file = True,
         mandatory = True,
     ),
+    "python_version": attr.string(
+        doc = """Whether to build this target and its transitive deps for a specific python version.
+
+Note that setting this attribute alone will not be enough as the python toolchain for the desired version
+also needs to be registered in the WORKSPACE or MODULE.bazel file.
+
+When using WORKSPACE, this may look like this,
+
+```
+load("@rules_python//python:repositories.bzl", "py_repositories", "python_register_toolchains")
+
+python_register_toolchains(
+    name = "python_toolchain_3_8",
+    python_version = "3.8.12",
+    # setting set_python_version_constraint makes it so that only matches py_* rule  
+    # which has this exact version set in the `python_version` attribute.
+    set_python_version_constraint = True,
+)
+
+# It's important to register the default toolchain last it will match any py_* target. 
+python_register_toolchains(
+    name = "python_toolchain",
+    python_version = "3.9",
+)
+```
+
+Configuring for MODULE.bazel may look like this:
+
+```
+python = use_extension("@rules_python//python/extensions:python.bzl", "python")
+python.toolchain(python_version = "3.8.12", is_default = False)
+python.toolchain(python_version = "3.9", is_default = True)
+```
+"""
+    ),
     "_run_tmpl": attr.label(
         allow_single_file = True,
         default = "//py/private:run.tmpl.sh",
@@ -150,17 +185,33 @@ _attrs = dict({
     "_interpreter_version_flag": attr.label(
         default = "//py:interpreter_version",
     ),
+    # Required for py_version attribute
+    "_allowlist_function_transition": attr.label(
+        default = "@bazel_tools//tools/allowlists/function_transition_allowlist",
+    ),
 })
 
 _attrs.update(**_py_library.attrs)
 
+def _python_version_transition_impl(_, attr):
+    if not attr.python_version:
+        return {}
+    return {"@rules_python//python/config_settings:python_version": str(attr.python_version)}
+
+_python_version_transition = transition(
+    implementation = _python_version_transition_impl,
+    inputs = [],
+    outputs = ["@rules_python//python/config_settings:python_version"],
+)
+
 py_base = struct(
     implementation = _py_binary_rule_impl,
     attrs = _attrs,
     toolchains = [
         PY_TOOLCHAIN,
         VENV_TOOLCHAIN,
     ],
+    cfg = _python_version_transition
 )
 
 py_binary = rule(
@@ -169,6 +220,7 @@ py_binary = rule(
     attrs = py_base.attrs,
     toolchains = py_base.toolchains,
     executable = True,
+    cfg = py_base.cfg
 )
 
 py_test = rule(
@@ -177,4 +229,5 @@ py_test = rule(
     attrs = py_base.attrs,
     toolchains = py_base.toolchains,
     test = True,
+    cfg = py_base.cfg
 )