fix: Allow PYTHONSTARTUP to define variables

With the current `//python/bin:repl` implementation, any variables
defined in `PYTHONSTARTUP` are not actually available in the REPL
itself. I accidentally omitted this in the first patch.

This patch fixes the issue and adds appropriate tests.

Author: philsc

python/bin/repl_stub.py

@@ -13,6 +13,9 @@
 The logic for PYTHONSTARTUP is handled in python/private/repl_template.py.
 """
 
+# Capture the globals from PYTHONSTARTUP so we can pass them on to the console.
+console_locals = globals().copy()
+
 import code
 import sys
 
@@ -26,4 +29,4 @@
     sys.ps2 = ""
 
 # We set the banner to an empty string because the repl_template.py file already prints the banner.
-code.interact(banner="", exitmsg=exitmsg)
+code.interact(local=console_locals, banner="", exitmsg=exitmsg)

python/private/repl_template.py

@@ -14,6 +14,8 @@ def start_repl():
         cprt = 'Type "help", "copyright", "credits" or "license" for more information.'
         sys.stderr.write("Python %s on %s\n%s\n" % (sys.version, sys.platform, cprt))
 
+    new_globals = {}
+
     # 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.
@@ -27,10 +29,10 @@ def start_repl():
             print(f"{type(error).__name__}: {error}")
         else:
             compiled_code = compile(source_code, filename=startup_file, mode="exec")
-            eval(compiled_code, {})
+            eval(compiled_code, new_globals)
 
     bazel_runfiles = runfiles.Create()
-    runpy.run_path(bazel_runfiles.Rlocation(STUB_PATH), run_name="__main__")
+    runpy.run_path(bazel_runfiles.Rlocation(STUB_PATH), init_globals=new_globals, run_name="__main__")
 
 
 if __name__ == "__main__":

tests/repl/repl_test.py

@@ -1,8 +1,10 @@
 import os
 import subprocess
 import sys
+import tempfile
 import unittest
 from typing import Iterable
+from pathlib import Path
 
 from python import runfiles
 
@@ -13,18 +15,25 @@
 EXPECT_TEST_MODULE_IMPORTABLE = os.environ["EXPECT_TEST_MODULE_IMPORTABLE"] == "1"
 
 
+# An arbitrary piece of code that sets some kind of variable. The variable needs to persist into the
+# actual shell.
+PYTHONSTARTUP_SETS_VAR = """\
+foo = 1234
+"""
+
 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:
+    def run_code_in_repl(self, lines: Iterable[str], *, env=None) -> 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),
+            env=env,
         ).strip()
 
     def test_repl_version(self):
@@ -69,6 +78,40 @@ def test_import_test_module_failure(self):
         )
         self.assertIn("ModuleNotFoundError: No module named 'test_module'", result)
 
+    def test_pythonstartup_gets_executed(self):
+        """Validates that we can use the variables from PYTHONSTARTUP in the console itself."""
+        with tempfile.TemporaryDirectory() as tempdir:
+            pythonstartup = Path(tempdir) / "pythonstartup.py"
+            pythonstartup.write_text(PYTHONSTARTUP_SETS_VAR)
+
+            env = os.environ.copy()
+            env["PYTHONSTARTUP"] = str(pythonstartup)
+
+            result = self.run_code_in_repl([
+                "print(f'The value of foo is {foo}')",
+            ], env=env)
+
+        self.assertIn("The value of foo is 1234", result)
+
+    def test_pythonstartup_doesnt_leak(self):
+        """Validates that we don't accidentally leak code into the console.
+
+        This test validates that a few of the variables we use in the template and stub are not
+        accessible in the REPL itself.
+        """
+        with tempfile.TemporaryDirectory() as tempdir:
+            pythonstartup = Path(tempdir) / "pythonstartup.py"
+            pythonstartup.write_text(PYTHONSTARTUP_SETS_VAR)
+
+            env = os.environ.copy()
+            env["PYTHONSTARTUP"] = str(pythonstartup)
+
+            for var_name in ("exitmsg", "sys", "code", "bazel_runfiles", "STUB_PATH"):
+                with self.subTest(var_name=var_name):
+                    result = self.run_code_in_repl([f"print({var_name})"], env=env)
+                    self.assertIn(f"NameError: name '{var_name}' is not defined", result)
+
+
 
 if __name__ == "__main__":
     unittest.main()