feat: Add support for REPLs

python/bin/BUILD.bazel

@@ -1,4 +1,5 @@
 load("//python/private:interpreter.bzl", _interpreter_binary = "interpreter_binary")
+load("//python/private:repl.bzl", "py_repl_binary")
 
 filegroup(
     name = "distribution",
@@ -22,3 +23,35 @@ label_flag(
     name = "python_src",
     build_setting_default = "//python:none",
 )
+
+py_repl_binary(
+    name = "repl",
+    stub = ":repl_stub",
+    visibility = ["//visibility:public"],
+    deps = [
+        ":repl_dep",
+        ":repl_stub_dep",
+    ],
+)
+
+# The user can replace this with their own stub. E.g. they can use this to
+# import ipython instead of the default shell.
+label_flag(
+    name = "repl_stub",
+    build_setting_default = "repl_stub.py",
+)
+
+# The user can modify this flag to make an interpreter shell library available
+# for the stub. E.g. if they switch the stub for an ipython-based one, then they
+# can point this at their version of ipython.
+label_flag(
+    name = "repl_stub_dep",
+    build_setting_default = "//python/private:empty",
+)
+
+# The user can modify this flag to make arbitrary PyInfo targets available for
+# import on the REPL.
+label_flag(
+    name = "repl_dep",
+    build_setting_default = "//python/private:empty",
+)

python/bin/repl_stub.py

@@ -0,0 +1,18 @@
+"""Simulates the REPL that Python spawns when invoking the binary with no arguments.
+
+The code module is responsible for the default shell.
+
+The import and `ocde.interact()` call here his is equivalent to doing:
+
+    $ python3 -m code
+    Python 3.11.2 (main, Mar 13 2023, 12:18:29) [GCC 12.2.0] on linux
+    Type "help", "copyright", "credits" or "license" for more information.
+    (InteractiveConsole)
+    >>>
+
+The logic for PYTHONSTARTUP is handled in python/private/repl_template.py.
+"""
+
+import code
+
+code.interact()

python/private/BUILD.bazel

@@ -817,6 +817,10 @@ current_interpreter_executable(
     visibility = ["//visibility:public"],
 )
 
+py_library(
+    name = "empty",
+)
+
 sentinel(
     name = "sentinel",
 )

python/private/repl.bzl

@@ -0,0 +1,82 @@
+"""Implementation of the rules to expose a REPL."""
+
+load("//python:py_binary.bzl", _py_binary = "py_binary")
+
+def _generate_repl_main_impl(ctx):
+    stub_repo = ctx.attr.stub.label.repo_name or ctx.workspace_name
+    stub_path = "/".join([stub_repo, ctx.file.stub.short_path])
+
+    # Point the generated main file at the stub.
+    ctx.actions.expand_template(
+        template = ctx.file._template,
+        output = ctx.outputs.out,
+        substitutions = {
+            "%stub_path%": stub_path,
+        },
+    )
+
+_generate_repl_main = rule(
+    implementation = _generate_repl_main_impl,
+    attrs = {
+        "out": attr.output(
+            mandatory = True,
+            doc = "The path to the file to generate.",
+        ),
+        "stub": attr.label(
+            mandatory = True,
+            allow_single_file = True,
+            doc = ("The stub responsible for actually invoking the final shell. " +
+                   "See the \"Customizing the REPL\" docs for details."),
+        ),
+        "_template": attr.label(
+            default = "//python/private:repl_template.py",
+            allow_single_file = True,
+            doc = "The template to use for generating `out`.",
+        ),
+    },
+    doc = """\
+Generates a "main" script for a py_binary target that starts a Python REPL.
+
+The template is designed to take care of the majority of the logic. The user
+customizes the exact shell that will be started via the stub. The stub is a
+simple shell script that imports the desired shell and then executes it.
+""",
+)
+
+def py_repl_binary(name, stub, deps = [], data = [], **kwargs):
+    """A py_binary target that executes a REPL when run.
+
+    The stub is the script that ultimately decides which shell the REPL will run.
+    It can be as simple as this:
+
+        import code
+        code.interact()
+
+    Or it can load something like IPython instead.
+
+    Args:
+        name: Name of the generated py_binary target.
+        stub: The script that invokes the shell.
+        deps: The dependencies of the py_binary.
+        data: The runtime dependencies of the py_binary.
+        **kwargs: Forwarded to the py_binary.
+    """
+    _generate_repl_main(
+        name = "%s_py" % name,
+        stub = stub,
+        out = "%s.py" % name,
+    )
+
+    _py_binary(
+        name = name,
+        srcs = [
+            ":%s.py" % name,
+        ],
+        data = data + [
+            stub,
+        ],
+        deps = deps + [
+            "//python/runfiles",
+        ],
+        **kwargs
+    )

python/private/repl_template.py

@@ -0,0 +1,31 @@
+import os
+import runpy
+from pathlib import Path
+
+from python.runfiles import runfiles
+
+STUB_PATH = "%stub_path%"
+
+
+def start_repl():
+    # Simulate Python's behavior when a valid startup script is defined by the
+    # PYTHONSTARTUP variable. If this file path fails to load, print the error
+    # and revert to the default behavior.
+    #
+    # See upstream for more information:
+    # https://docs.python.org/3/using/cmdline.html#envvar-PYTHONSTARTUP
+    if startup_file := os.getenv("PYTHONSTARTUP"):
+        try:
+            source_code = Path(startup_file).read_text()
+        except Exception as error:
+            print(f"{type(error).__name__}: {error}")
+        else:
+            compiled_code = compile(source_code, filename=startup_file, mode="exec")
+            eval(compiled_code, {})
+
+    bazel_runfiles = runfiles.Create()
+    runpy.run_path(bazel_runfiles.Rlocation(STUB_PATH), run_name="__main__")
+
+
+if __name__ == "__main__":
+    start_repl()

tests/repl/BUILD.bazel

@@ -0,0 +1,44 @@
+load("//python:py_library.bzl", "py_library")
+load("//tests/support:sh_py_run_test.bzl", "py_reconfig_test")
+
+# A library that adds a special import path only when this is specified as a
+# dependency. This makes it easy for a dependency to have this import path
+# available without the top-level target being able to import the module.
+py_library(
+    name = "helper/test_module",
+    srcs = [
+        "helper/test_module.py",
+    ],
+    imports = [
+        "helper",
+    ],
+)
+
+py_reconfig_test(
+    name = "repl_without_dep_test",
+    srcs = ["repl_test.py"],
+    data = [
+        "//python/bin:repl",
+    ],
+    env = {
+        # The helper/test_module should _not_ be importable for this test.
+        "EXPECT_TEST_MODULE_IMPORTABLE": "0",
+    },
+    main = "repl_test.py",
+    python_version = "3.12",
+)
+
+py_reconfig_test(
+    name = "repl_with_dep_test",
+    srcs = ["repl_test.py"],
+    data = [
+        "//python/bin:repl",
+    ],
+    env = {
+        # The helper/test_module _should_ be importable for this test.
+        "EXPECT_TEST_MODULE_IMPORTABLE": "1",
+    },
+    main = "repl_test.py",
+    python_version = "3.12",
+    repl_dep = ":helper/test_module",
+)

tests/repl/helper/test_module.py

@@ -0,0 +1,5 @@
+"""This is a file purely intended for validating //python/bin:repl."""
+
+
+def print_hello():
+    print("Hello World")

tests/repl/repl_test.py

@@ -0,0 +1,74 @@
+import os
+import subprocess
+import sys
+import unittest
+from typing import Iterable
+
+from python import runfiles
+
+rfiles = runfiles.Create()
+
+# Signals the tests below whether we should be expecting the import of
+# helpers/test_module.py on the REPL to work or not.
+EXPECT_TEST_MODULE_IMPORTABLE = os.environ["EXPECT_TEST_MODULE_IMPORTABLE"] == "1"
+
+
+class ReplTest(unittest.TestCase):
+    def setUp(self):
+        self.repl = rfiles.Rlocation("rules_python/python/bin/repl")
+        assert self.repl
+
+    def run_code_in_repl(self, lines: Iterable[str]) -> str:
+        """Runs the lines of code in the REPL and returns the text output."""
+        return subprocess.check_output(
+            [self.repl],
+            text=True,
+            stderr=subprocess.STDOUT,
+            input="\n".join(lines),
+        ).strip()
+
+    def test_repl_version(self):
+        """Validates that we can successfully execute arbitrary code on the REPL."""
+
+        result = self.run_code_in_repl(
+            [
+                "import sys",
+                "v = sys.version_info",
+                "print(f'version: {v.major}.{v.minor}')",
+            ]
+        )
+        self.assertIn("version: 3.12", result)
+
+    def test_cannot_import_test_module_directly(self):
+        """Validates that we cannot import helper/test_module.py since it's not a direct dep."""
+        with self.assertRaises(ModuleNotFoundError):
+            import test_module
+
+    @unittest.skipIf(
+        not EXPECT_TEST_MODULE_IMPORTABLE, "test only works without repl_dep set"
+    )
+    def test_import_test_module_success(self):
+        """Validates that we can import helper/test_module.py when repl_dep is set."""
+        result = self.run_code_in_repl(
+            [
+                "import test_module",
+                "test_module.print_hello()",
+            ]
+        )
+        self.assertIn("Hello World", result)
+
+    @unittest.skipIf(
+        EXPECT_TEST_MODULE_IMPORTABLE, "test only works without repl_dep set"
+    )
+    def test_import_test_module_failure(self):
+        """Validates that we cannot import helper/test_module.py when repl_dep isn't set."""
+        result = self.run_code_in_repl(
+            [
+                "import test_module",
+            ]
+        )
+        self.assertIn("ModuleNotFoundError: No module named 'test_module'", result)
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/support/sh_py_run_test.bzl

@@ -38,6 +38,8 @@ def _perform_transition_impl(input_settings, attr, base_impl):
         settings["//command_line_option:extra_toolchains"] = attr.extra_toolchains
     if attr.python_src:
         settings["//python/bin:python_src"] = attr.python_src
+    if attr.repl_dep:
+        settings["//python/bin:repl_dep"] = attr.repl_dep
     if attr.venvs_use_declare_symlink:
         settings["//python/config_settings:venvs_use_declare_symlink"] = attr.venvs_use_declare_symlink
     if attr.venvs_site_packages:
@@ -47,6 +49,7 @@ def _perform_transition_impl(input_settings, attr, base_impl):
 _RECONFIG_INPUTS = [
     "//python/config_settings:bootstrap_impl",
     "//python/bin:python_src",
+    "//python/bin:repl_dep",
     "//command_line_option:extra_toolchains",
     "//python/config_settings:venvs_use_declare_symlink",
     "//python/config_settings:venvs_site_packages",
@@ -70,6 +73,7 @@ toolchain.
 """,
     ),
     "python_src": attrb.Label(),
+    "repl_dep": attrb.Label(),
     "venvs_site_packages": attrb.String(),
     "venvs_use_declare_symlink": attrb.String(),
 }