Seal Python with a VirtualEnv toolchain

BUILD

load("@rules_python//python:defs.bzl", "py_runtime", "py_runtime_pair")
load(":venv.bzl", "py_venv")
load(":wrapper.bzl", "py_wrapper")

py_venv(
    name = "py3_venv",
    # caching the venv interpreter doesn't work for two reasons:
    #
    # * there are no platform dependencies on this target: the cache will contain binaries for the wrong OS
    # * there are no dependencies on the interpreter binary or standard library: the cache may contain stale versions
    #
    # fixing this properly a bit of a pain...
    # the preferred solution is to use a platform specific tarball that does not need to be relocated.
    # unfortunately the Python organization only builds/hosts _installers_ and does not provide binary tarballs.
    # building/maintaining our own binary tarballs is quite a bit of work, if we wait someone else may do it for us.
    # the second best would to add the interpreter and standard library as inputs to the py_venv rule.
    # this is not a terrible interim solution, it just needs a repository_rule and a python script to hash
    # said interpreter and libraries, but it's not free and does not provide hermeticicity (binary tarballs would).
    #
    # "no-remote-cache" fixes bullet 1. "no-cache" fixes bullet 1 and mostly avoids bullet 2.
    tags = ["no-cache"],
)

py_wrapper(
    name = "py3_wrapper",
    venv = ":py3_venv",
)

# python2 doesn't have proper venv support so invoke it with -S and hope for the best
py_runtime(
    name = "py2_runtime",
    interpreter = ":py2wrapper.sh",
    python_version = "PY2",
)

# python3 gets a real venv, this is required to run tensorflow due to expectations that
# site.getsitepackages() returns a list of directories
py_runtime(
    name = "py3_runtime",
    files = [":py3_venv"],  # annoying: https://github.com/bazelbuild/bazel/issues/4286
    interpreter = ":py3_wrapper",
    python_version = "PY3",
    stub_shebang = "#!/usr/bin/env python3",
)

py_runtime_pair(
    name = "py23_runtime_pair",
    py2_runtime = ":py2_runtime",
    py3_runtime = ":py3_runtime",
)

toolchain(
    name = "python",
    toolchain = ":py23_runtime_pair",
    toolchain_type = "@bazel_tools//tools/python:toolchain_type",
)

create_venv.sh

#!/usr/bin/env sh
set -e
"$1" -m venv --without-pip "$2"

# remove unused activation scripts that break caching
rm "$2/bin/"[Aa]ctivate*

LICENSE

Copyright (c) 2022 Nathan Howell
Copyright (C) 2022 Embedl AB

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.

py3wrapper.sh

#!/usr/bin/env bash
ROOT=$(dirname "$0")
# disable pycache which produces non-deterministic outputs
export PYTHONDONTWRITEBYTECODE=1
exec "$ROOT/venv/bin/python3" -B "$@"

venv.bzl

def _impl(ctx):
    version = ctx.attr.version.split(".", 2)
    if len(version) != 2:
        fail("Expected version in major.minor format, e.g. 3.7. Found: {}".format(ctx.attr.version))

    py_minor = ctx.actions.declare_file("venv/bin/python{}".format(ctx.attr.version))
    py_major = ctx.actions.declare_file("venv/bin/python{}".format(version[0]))
    python = ctx.actions.declare_file("venv/bin/python")
    venv_cfg = ctx.actions.declare_file("venv/pyvenv.cfg")

    ctx.actions.run(
        executable = ctx.executable._create_venv,
        execution_requirements = {tag: "true" for tag in ctx.attr.tags},
        arguments = ["python{}".format(ctx.attr.version), venv_cfg.dirname],
        outputs = [python, py_major, py_minor, venv_cfg],
        use_default_shell_env = True,
    )

    return [
        DefaultInfo(
            executable = py_major,
            files = depset([python, py_major, py_minor, venv_cfg]),
            runfiles = ctx.runfiles(
                files = [python, py_major, py_minor, venv_cfg],
            ),
        ),
    ]

py_venv = rule(
    implementation = _impl,
    executable = True,
    attrs = {
        "version": attr.string(default = "3.7"),
        "_create_venv": attr.label(
            default = ":create_venv.sh",
            allow_single_file = True,
            executable = True,
            cfg = "target",
        ),
    },
)

WORKSPACE

# register local toolchains before loading other workspaces
register_toolchains("//toolchains:python")

wrapper.bzl

def _impl(ctx):
    wrapper = ctx.actions.declare_file("py3wrapper.sh")

    # we are required to create a new file when returning an executable...
    # so even though no substitutions are being made this seems about the same
    # as running `cp`
    ctx.actions.expand_template(
        template = ctx.file._template,
        output = wrapper,
        is_executable = True,
        substitutions = {},
    )

    return [
        DefaultInfo(
            executable = wrapper,
            runfiles = ctx.runfiles(
                files = [wrapper],
                transitive_files = depset(ctx.files.venv),
            ),
        ),
    ]

py_wrapper = rule(
    implementation = _impl,
    executable = True,
    attrs = {
        "venv": attr.label(
            mandatory = True,
            executable = True,
            cfg = "target",
        ),
        "_template": attr.label(
            default = ":py3wrapper.sh",
            allow_single_file = True,
        ),
    },
)