Allow PyInfo in mojo_binary / mojo_test deps (#3)

This enables python interop in the way you'd expect. I don't know if we
need this on `mojo_library` ever. The key here was to make sure all
DefaultInfo.default_runfiles from the python targets was included,
because PyInfo.transitive_files doesn't include native libraries that
the dependency vendors

Author: keith

.bazelrc

@@ -0,0 +1,2 @@
+common --incompatible_strict_action_env
+common --test_output=errors

MODULE.bazel

@@ -7,9 +7,27 @@ module(
 bazel_dep(name = "bazel_skylib", version = "1.7.1")
 bazel_dep(name = "platforms", version = "0.0.11")
 bazel_dep(name = "rules_cc", version = "0.1.1")
+bazel_dep(name = "rules_python", version = "1.0.0")
 
 mojo = use_extension("//mojo:extensions.bzl", "mojo")
 mojo.toolchain()
 use_repo(mojo, "mojo_toolchains")
 
 register_toolchains("@mojo_toolchains//...")
+
+_DEFAULT_PYTHON_VERSION = "3.12"
+
+python = use_extension("@rules_python//python/extensions:python.bzl", "python", dev_dependency = True)
+python.toolchain(
+    ignore_root_user_error = True,
+    is_default = True,
+    python_version = _DEFAULT_PYTHON_VERSION,
+)
+
+pip = use_extension("@rules_python//python/extensions:pip.bzl", "pip", dev_dependency = True)
+pip.parse(
+    hub_name = "rules_mojo_test_deps",
+    python_version = _DEFAULT_PYTHON_VERSION,
+    requirements_lock = "tests/python/requirements.txt",
+)
+use_repo(pip, "rules_mojo_test_deps")

mojo/private/mojo_binary_test.bzl

@@ -3,6 +3,7 @@
 load("@bazel_skylib//lib:paths.bzl", "paths")
 load("@bazel_skylib//rules:common_settings.bzl", "BuildSettingInfo")
 load("@bazel_tools//tools/cpp:toolchain_utils.bzl", "find_cpp_toolchain", "use_cpp_toolchain")
+load("@rules_python//python:py_info.bzl", "PyInfo")
 load("//mojo:providers.bzl", "MojoInfo")
 load(":utils.bzl", "MOJO_EXTENSIONS", "collect_mojoinfo")
 
@@ -16,7 +17,7 @@ _ATTRS = {
     ),
     "copts": attr.string_list(),
     "deps": attr.label_list(
-        providers = [[CcInfo], [MojoInfo]],
+        providers = [[CcInfo], [MojoInfo], [PyInfo]],
     ),
     "data": attr.label_list(allow_files = True),
     "enable_assertions": attr.bool(default = True),
@@ -28,6 +29,7 @@ _ATTRS = {
 
 _TOOLCHAINS = use_cpp_toolchain() + [
     "//:toolchain_type",
+    "@bazel_tools//tools/python:toolchain_type",
 ]
 
 def _find_main(name, srcs, main):
@@ -56,8 +58,9 @@ def _find_main(name, srcs, main):
     fail("Multiple Mojo files provided, but no main file specified. Please set 'main = \"foo.mojo\"' to disambiguate.")
 
 def _mojo_binary_test_implementation(ctx):
-    mojo_toolchain = ctx.toolchains["//:toolchain_type"].mojo_toolchain_info
     cc_toolchain = find_cpp_toolchain(ctx)
+    mojo_toolchain = ctx.toolchains["//:toolchain_type"].mojo_toolchain_info
+    py_toolchain = ctx.toolchains["@bazel_tools//tools/python:toolchain_type"]
 
     object_file = ctx.actions.declare_file(ctx.label.name + ".lo")
     args = ctx.actions.args()
@@ -138,20 +141,53 @@ def _mojo_binary_test_implementation(ctx):
 
     data = ctx.attr.data
     runfiles = ctx.runfiles(ctx.files.data)
-    transitive_runfiles = []
+    transitive_runfiles = [
+        ctx.runfiles(transitive_files = py_toolchain.py3_runtime.files),
+    ]
     for target in data:
         transitive_runfiles.append(target[DefaultInfo].default_runfiles)
 
     # Collect transitive shared libraries that must exist at runtime
+    python_imports = []
     for target in ctx.attr.deps + mojo_toolchain.implicit_deps:
+        transitive_runfiles.append(target[DefaultInfo].default_runfiles)
+
+        if PyInfo in target:
+            python_imports.append(target[PyInfo].imports)
+            transitive_runfiles.append(
+                ctx.runfiles(transitive_files = target[PyInfo].transitive_sources),
+            )
+
         if CcInfo not in target:
             continue
         for linker_input in target[CcInfo].linking_context.linker_inputs.to_list():
             for library in linker_input.libraries:
                 if library.dynamic_library and not library.pic_static_library and not library.static_library:
                     transitive_runfiles.append(ctx.runfiles(transitive_files = depset([library.dynamic_library])))
 
-    runtime_env = dict(ctx.attr.env)
+    python_path = ""
+    for path in depset(transitive = python_imports).to_list():
+        python_path += "../" + path + ":"
+
+    # https://github.com/bazelbuild/rules_python/issues/2262
+    libpython = None
+    for file in py_toolchain.py3_runtime.files.to_list():
+        if file.basename.startswith("libpython"):
+            libpython = file.short_path
+            break  # if there are multiple any of them should work and they are likely symlinks to each other
+
+    if not libpython:
+        fail("failed to find libpython, please report this at https://github.com/modular/rules_mojo/issues")
+
+    runtime_env = dict(ctx.attr.env) | {
+        "MODULAR_PYTHON_EXECUTABLE": py_toolchain.py3_runtime.interpreter.short_path,
+        "MOJO_PYTHON": py_toolchain.py3_runtime.interpreter.short_path,
+        "MOJO_PYTHON_LIBRARY": libpython,
+        "PYTHONEXECUTABLE": py_toolchain.py3_runtime.interpreter.short_path,
+        "PYTHONNOUSERSITE": "affirmative",
+        "PYTHONPATH": python_path,
+        "PYTHONSAFEPATH": "affirmative",
+    }
     for key, value in runtime_env.items():
         runtime_env[key] = ctx.expand_make_variables(
             "env",

tests/python/BUILD.bazel

@@ -0,0 +1,15 @@
+load("@rules_mojo_test_deps//:requirements.bzl", "requirement")
+load("//mojo:mojo_test.bzl", "mojo_test")
+
+mojo_test(
+    name = "basic_python_test",
+    srcs = ["basic_python_test.mojo"],
+)
+
+mojo_test(
+    name = "deps_python_test",
+    srcs = ["deps_python_test.mojo"],
+    deps = [
+        requirement("numpy"),
+    ],
+)

tests/python/basic_python_test.mojo

@@ -0,0 +1,7 @@
+from python import Python
+from testing import assert_true
+
+def main():
+    sys = Python.import_module("sys")
+    print("Python executable:", sys.executable)
+    assert_true(not sys.executable.startswith("/usr"))

tests/python/deps_python_test.mojo

@@ -0,0 +1,18 @@
+from python import Python
+from testing import assert_equal
+import os
+import subprocess
+
+def test_basic_numpy_example():
+    var np = Python.import_module("numpy")
+    var array = np.array(
+        Python.list(
+            Python.list(1, 2, 3),
+            Python.list(4, 5, 6)
+        )
+    )
+    assert_equal(array.shape, Python.tuple(2, 3))
+
+
+def main():
+    test_basic_numpy_example()

tests/python/requirements.txt

@@ -0,0 +1,34 @@
+# This file was autogenerated by uv via the following command:
+#    uv export -o requirements.txt
+numpy==2.2.5 \
+    --hash=sha256:02f226baeefa68f7d579e213d0f3493496397d8f1cff5e2b222af274c86a552a \
+    --hash=sha256:059b51b658f4414fff78c6d7b1b4e18283ab5fa56d270ff212d5ba0c561846f4 \
+    --hash=sha256:0cd48122a6b7eab8f06404805b1bd5856200e3ed6f8a1b9a194f9d9054631beb \
+    --hash=sha256:1a161c2c79ab30fe4501d5a2bbfe8b162490757cf90b7f05be8b80bc02f7bb8e \
+    --hash=sha256:261a1ef047751bb02f29dfe337230b5882b54521ca121fc7f62668133cb119c9 \
+    --hash=sha256:2ba321813a00e508d5421104464510cc962a6f791aa2fca1c97b1e65027da80d \
+    --hash=sha256:2c1a1c6ccce4022383583a6ded7bbcda22fc635eb4eb1e0a053336425ed36dfa \
+    --hash=sha256:352d330048c055ea6db701130abc48a21bec690a8d38f8284e00fab256dc1376 \
+    --hash=sha256:3a801fef99668f309b88640e28d261991bfad9617c27beda4a3aec4f217ea073 \
+    --hash=sha256:3d14b17b9be5f9c9301f43d2e2a4886a33b53f4e6fdf9ca2f4cc60aeeee76372 \
+    --hash=sha256:4520caa3807c1ceb005d125a75e715567806fed67e315cea619d5ec6e75a4191 \
+    --hash=sha256:47834cde750d3c9f4e52c6ca28a7361859fcaf52695c7dc3cc1a720b8922683e \
+    --hash=sha256:47f9ed103af0bc63182609044b0490747e03bd20a67e391192dde119bf43d52f \
+    --hash=sha256:54088a5a147ab71a8e7fdfd8c3601972751ded0739c6b696ad9cb0343e21ab73 \
+    --hash=sha256:55f09e00d4dccd76b179c0f18a44f041e5332fd0e022886ba1c0bbf3ea4a18d0 \
+    --hash=sha256:5a0ac90e46fdb5649ab6369d1ab6104bfe5854ab19b645bf5cda0127a13034ae \
+    --hash=sha256:8b4c0773b6ada798f51f0f8e30c054d32304ccc6e9c5d93d46cb26f3d385ab19 \
+    --hash=sha256:8dfa94b6a4374e7851bbb6f35e6ded2120b752b063e6acdd3157e4d2bb922eba \
+    --hash=sha256:97c8425d4e26437e65e1d189d22dff4a079b747ff9c2788057bfb8114ce1e133 \
+    --hash=sha256:9d75f338f5f79ee23548b03d801d28a505198297534f62416391857ea0479571 \
+    --hash=sha256:a4cbdef3ddf777423060c6f81b5694bad2dc9675f110c4b2a60dc0181543fac7 \
+    --hash=sha256:a9c0d994680cd991b1cb772e8b297340085466a6fe964bc9d4e80f5e2f43c291 \
+    --hash=sha256:abe38cd8381245a7f49967a6010e77dbf3680bd3627c0fe4362dd693b404c7f8 \
+    --hash=sha256:c26843fd58f65da9491165072da2cccc372530681de481ef670dcc8e27cfb066 \
+    --hash=sha256:c8b82a55ef86a2d8e81b63da85e55f5537d2157165be1cb2ce7cfa57b6aef38b \
+    --hash=sha256:ced69262a8278547e63409b2653b372bf4baff0870c57efa76c5703fd6543282 \
+    --hash=sha256:d403c84991b5ad291d3809bace5e85f4bbf44a04bdc9a88ed2bb1807b3360bb8 \
+    --hash=sha256:d8882a829fd779f0f43998e931c466802a77ca1ee0fe25a3abe50278616b1471 \
+    --hash=sha256:e8b025c351b9f0e8b5436cf28a07fa4ac0204d67b38f01433ac7f9b870fa38c6 \
+    --hash=sha256:ec31367fd6a255dc8de4772bd1658c3e926d8e860a0b6e922b615e532d320ddc \
+    --hash=sha256:ee461a4eaab4f165b68780a6a1af95fb23a29932be7569b9fab666c407969051