Fix issues with REPL implementation (pantsbuild#10597)

There were a few issues with our REPL support:

1. We ran the REPL in a tmpdir under the repo root, with the cwd was the repo root.
So if the repo root also happened to also be a source root we'd import from the
original sources in the repo, instead of from the chrooted copy.
2. There was custom target filtering logic on the repl implementations, that didn't handle
codegen (and was redundant anyway).

This change fixes these by:

1. Running the REPL in a chroot not under the repo root.
2. Getting rid of the custom target filtering. Our python/ipython REPL implementations
just use the standard mechanisms for this.

This change also refactors out a requirements_pex_from_targets_request helper function
that creates a PexFromTargetsRequest that can generate a pex containing just requirements.
We now use this in pytest_runner.py and in repl.py. This increases the chances of getting
cache hits, as this helper ensures that, e.g., the pex output file name is uniform.

To further increase cache hits, this change resolves ipython in a separate pex, instead of
in the requirements pex.

As a result, we can now see generated code in python/ipython repls.

[ci skip-rust]

[ci skip-build-wheels]

src/python/pants/backend/python/rules/pex_from_targets.py

@@ -83,6 +83,28 @@ def __init__(
         self.additional_inputs = additional_inputs
         self.description = description
 
+    @classmethod
+    def for_requirements(
+        cls, addresses: Addresses, *, zip_safe: bool = False
+    ) -> "PexFromTargetsRequest":
+        """Create an instance that can be used to get a requirements pex.
+
+        Useful to ensure that these requests are uniform (e.g., the using the same output filename),
+        so that the underlying pexes are more likely to be reused instead of re-resolved.
+
+        We default to zip_safe=False because there are various issues with running zipped pexes
+        directly, and it's best to only use those if you're sure it's the right thing to do.
+        Also, pytest must use zip_safe=False for performance reasons (see comment in
+        pytest_runner.py) and we get more re-use of pexes if other uses follow suit.
+        This default is a helpful nudge in that direction.
+        """
+        return PexFromTargetsRequest(
+            addresses=sorted(addresses),
+            output_filename="requirements.pex",
+            include_source_files=False,
+            additional_args=() if zip_safe else ("--not-zip-safe",),
+        )
+
 
 @dataclass(frozen=True)
 class TwoStepPexFromTargetsRequest:

src/python/pants/backend/python/rules/pytest_runner.py

@@ -1,7 +1,6 @@
 # Copyright 2018 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
 
-import functools
 import itertools
 import logging
 from dataclasses import dataclass
@@ -89,11 +88,6 @@ async def setup_pytest_for_target(
         python_setup,
     )
 
-    # Ensure all pexes we merge via PEX_PATH to form the test runner use the interpreter constraints
-    # of the tests. This is handled by CreatePexFromTargetClosure, but we must pass this through for
-    # CreatePex requests.
-    pex_request = functools.partial(PexRequest, interpreter_constraints=interpreter_constraints)
-
     # NB: We set `--not-zip-safe` because Pytest plugin discovery, which uses
     # `importlib_metadata` and thus `zipp`, does not play nicely when doing import magic directly
     # from zip files. `zipp` has pathologically bad behavior with large zipfiles.
@@ -104,31 +98,25 @@ async def setup_pytest_for_target(
 
     pytest_pex_request = Get(
         Pex,
-        PexRequest,
-        pex_request(
+        PexRequest(
             output_filename="pytest.pex",
             requirements=PexRequirements(pytest.get_requirement_strings()),
+            interpreter_constraints=interpreter_constraints,
             additional_args=additional_args_for_pytest,
         ),
     )
 
+    # Defaults to zip_safe=False.
     requirements_pex_request = Get(
-        Pex,
-        PexFromTargetsRequest(
-            addresses=test_addresses,
-            output_filename="requirements.pex",
-            include_source_files=False,
-            additional_args=additional_args_for_pytest,
-        ),
+        Pex, PexFromTargetsRequest, PexFromTargetsRequest.for_requirements(test_addresses)
     )
 
     test_runner_pex_request = Get(
         Pex,
-        PexRequest,
-        pex_request(
+        PexRequest(
+            interpreter_constraints=interpreter_constraints,
             output_filename="test_runner.pex",
             entry_point="pytest:main",
-            interpreter_constraints=interpreter_constraints,
             additional_args=(
                 "--pex-path",
                 # TODO(John Sirois): Support shading python binaries:

src/python/pants/backend/python/rules/repl.py

@@ -1,70 +1,93 @@
 # Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
+from typing import Tuple
 
-from pants.backend.python.rules.pex import Pex
+from pants.backend.python.rules.pex import Pex, PexRequest, PexRequirements
 from pants.backend.python.rules.pex_from_targets import PexFromTargetsRequest
 from pants.backend.python.rules.python_sources import PythonSourceFiles, PythonSourceFilesRequest
 from pants.backend.python.subsystems.ipython import IPython
-from pants.backend.python.target_types import PythonSources
 from pants.core.goals.repl import ReplImplementation, ReplRequest
+from pants.engine.addresses import Addresses
 from pants.engine.fs import Digest, MergeDigests
 from pants.engine.rules import Get, MultiGet, collect_rules, rule
 from pants.engine.unions import UnionRule
 
 
 class PythonRepl(ReplImplementation):
     name = "python"
-    required_fields = (PythonSources,)
 
 
 @rule
 async def create_python_repl_request(repl: PythonRepl) -> ReplRequest:
-    pex_request = Get(
+    requirements_request = Get(
         Pex,
-        PexFromTargetsRequest(
-            (tgt.address for tgt in repl.targets),
-            output_filename="python.pex",
-            include_source_files=False,
-        ),
+        PexFromTargetsRequest,
+        PexFromTargetsRequest.for_requirements(Addresses(tgt.address for tgt in repl.targets)),
+    )
+    sources_request = Get(
+        PythonSourceFiles, PythonSourceFilesRequest(repl.targets, include_files=True)
     )
-    sources_request = Get(PythonSourceFiles, PythonSourceFilesRequest(repl.targets))
-    pex, sources = await MultiGet(pex_request, sources_request)
+    requirements_pex, sources = await MultiGet(requirements_request, sources_request)
     merged_digest = await Get(
-        Digest, MergeDigests((pex.digest, sources.source_files.snapshot.digest))
+        Digest, MergeDigests((requirements_pex.digest, sources.source_files.snapshot.digest))
     )
+    chrooted_source_roots = [repl.in_chroot(sr) for sr in sources.source_roots]
     return ReplRequest(
         digest=merged_digest,
-        binary_name=pex.output_filename,
-        env={"PEX_EXTRA_SYS_PATH": ":".join(sources.source_roots)},
+        args=(repl.in_chroot(requirements_pex.output_filename),),
+        env={"PEX_EXTRA_SYS_PATH": ":".join(chrooted_source_roots)},
     )
 
 
 class IPythonRepl(ReplImplementation):
     name = "ipython"
-    required_fields = (PythonSources,)
 
 
 @rule
 async def create_ipython_repl_request(repl: IPythonRepl, ipython: IPython) -> ReplRequest:
-    pex_request = Get(
+    # Note that we get an intermediate PexRequest here (instead of going straight to a Pex)
+    # so that we can get the interpreter constraints for use in ipython_request.
+    requirements_pex_request = await Get(
+        PexRequest,
+        PexFromTargetsRequest,
+        PexFromTargetsRequest.for_requirements(Addresses(tgt.address for tgt in repl.targets)),
+    )
+
+    requirements_request = Get(Pex, PexRequest, requirements_pex_request)
+
+    sources_request = Get(
+        PythonSourceFiles, PythonSourceFilesRequest(repl.targets, include_files=True)
+    )
+
+    req_pex_path = repl.in_chroot(requirements_pex_request.output_filename)
+    ipython_request = Get(
         Pex,
-        PexFromTargetsRequest(
-            (tgt.address for tgt in repl.targets),
+        PexRequest(
             output_filename="ipython.pex",
             entry_point=ipython.entry_point,
-            additional_requirements=ipython.all_requirements,
-            include_source_files=True,
+            requirements=PexRequirements(ipython.all_requirements),
+            interpreter_constraints=requirements_pex_request.interpreter_constraints,
+            additional_args=("--pex-path", req_pex_path),
         ),
     )
-    sources_request = Get(PythonSourceFiles, PythonSourceFilesRequest(repl.targets))
-    pex, sources = await MultiGet(pex_request, sources_request)
+
+    requirements_pex, sources, ipython_pex = await MultiGet(
+        requirements_request, sources_request, ipython_request
+    )
     merged_digest = await Get(
-        Digest, MergeDigests((pex.digest, sources.source_files.snapshot.digest))
+        Digest,
+        MergeDigests(
+            (requirements_pex.digest, sources.source_files.snapshot.digest, ipython_pex.digest)
+        ),
     )
+    chrooted_source_roots = [repl.in_chroot(sr) for sr in sources.source_roots]
+    args: Tuple[str, ...] = (repl.in_chroot(ipython_pex.output_filename),)
+    if ipython.options.ignore_cwd:
+        args = args + ("--ignore-cwd",)
     return ReplRequest(
         digest=merged_digest,
-        binary_name=pex.output_filename,
-        env={"PEX_EXTRA_SYS_PATH": ":".join(sources.source_roots)},
+        args=args,
+        env={"PEX_EXTRA_SYS_PATH": ":".join(chrooted_source_roots)},
     )
 
 

src/python/pants/backend/python/subsystems/ipython.py

@@ -10,7 +10,23 @@ class IPython(PythonToolBase):
     """The IPython enhanced REPL (https://ipython.org/)."""
 
     options_scope = "ipython"
-    default_version = "ipython==5.8.0"
+    default_version = "ipython==7.16.1"  # The last version to support Python 3.6.
     default_extra_requirements: List[str] = []
     default_entry_point = "IPython:start_ipython"
-    default_interpreter_constraints = ["CPython>=2.7,<3", "CPython>=3.4"]
+    default_interpreter_constraints = ["CPython>=3.6"]
+
+    @classmethod
+    def register_options(cls, register):
+        super().register_options(register)
+        register(
+            "--ignore-cwd",
+            type=str,
+            advanced=True,
+            default=True,
+            help="Whether to tell IPython not to put the CWD on the import path. "
+            "Normally you want this to be True, so that imports come from the hermetic "
+            "environment Pants creates.  However IPython<7.13.0 doesn't support this option, "
+            "so if you're using an earlier version (e.g., because you have Python 2.7 code) "
+            "then you will need to set this to False, and you may have issues with imports "
+            "from your CWD shading the hermetic environment.",
+        )

src/python/pants/core/goals/BUILD

@@ -15,6 +15,12 @@ python_integration_tests(
     'examples/src/python/example:hello_directory',
     '//:isort_cfg',
     'examples:isort_cfg',
+    ':pydevd'
   ],
   timeout=240,
 )
+
+python_requirement_library(
+  name='pydevd',
+  requirements=[python_requirement('pydevd-pycharm')]
+)

src/python/pants/core/goals/repl.py

@@ -1,6 +1,6 @@
 # Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
 # Licensed under the Apache License, Version 2.0 (see LICENSE).
-
+import os
 from abc import ABC
 from dataclasses import dataclass
 from pathlib import PurePath
@@ -12,7 +12,7 @@
 from pants.engine.goal import Goal, GoalSubsystem
 from pants.engine.process import InteractiveProcess, InteractiveRunner
 from pants.engine.rules import Get, collect_rules, goal_rule
-from pants.engine.target import Field, Target, Targets, TransitiveTargets
+from pants.engine.target import Targets, TransitiveTargets
 from pants.engine.unions import UnionMembership, union
 from pants.option.global_options import GlobalOptions
 from pants.util.contextutil import temporary_dir
@@ -23,17 +23,18 @@
 @union
 @dataclass(frozen=True)
 class ReplImplementation(ABC):
-    """This type proxies from the top-level `repl` goal to a specific REPL implementation for a
-    specific language or languages."""
+    """A REPL implementation for a specific language or runtime.
 
-    name: ClassVar[str]
-    required_fields: ClassVar[Tuple[Type[Field], ...]]
+    Proxies from the top-level `repl` goal to an actual implementation.
+    """
 
+    name: ClassVar[str]
     targets: Targets
+    # Absolute path of the chroot the sources will be materialized to.
+    chroot: str
 
-    @classmethod
-    def is_valid(cls, tgt: Target) -> bool:
-        return tgt.has_fields(cls.required_fields)
+    def in_chroot(self, relpath: str) -> str:
+        return os.path.join(self.chroot, relpath)
 
 
 class ReplSubsystem(GoalSubsystem):
@@ -65,14 +66,14 @@ class Repl(Goal):
 @dataclass(unsafe_hash=True)
 class ReplRequest:
     digest: Digest
-    binary_name: str
+    args: Tuple[str, ...]
     env: FrozenDict[str, str]
 
     def __init__(
-        self, *, digest: Digest, binary_name: str, env: Optional[Mapping[str, str]] = None,
+        self, *, digest: Digest, args: Tuple[str, ...], env: Optional[Mapping[str, str]] = None,
     ) -> None:
         self.digest = digest
-        self.binary_name = binary_name
+        self.args = args
         self.env = FrozenDict(env or {})
 
 
@@ -87,6 +88,8 @@ async def run_repl(
     union_membership: UnionMembership,
     global_options: GlobalOptions,
 ) -> Repl:
+    # TODO: When we support multiple languages, detect the default repl to use based
+    #  on the targets.  For now we default to the python repl.
     repl_shell_name = repl_subsystem.shell or "python"
 
     implementations: Dict[str, Type[ReplImplementation]] = {
@@ -101,21 +104,17 @@ async def run_repl(
         )
         return Repl(-1)
 
-    repl_impl = repl_implementation_cls(
-        targets=Targets(
-            tgt for tgt in transitive_targets.closure if repl_implementation_cls.is_valid(tgt)
-        )
-    )
-    request = await Get(ReplRequest, ReplImplementation, repl_impl)
-
     with temporary_dir(root_dir=global_options.options.pants_workdir, cleanup=False) as tmpdir:
         tmpdir_relative_path = PurePath(tmpdir).relative_to(build_root.path).as_posix()
-        exe_path = PurePath(tmpdir, request.binary_name).as_posix()
+        repl_impl = repl_implementation_cls(
+            targets=Targets(transitive_targets.closure), chroot=tmpdir
+        )
+        request = await Get(ReplRequest, ReplImplementation, repl_impl)
+
         workspace.write_digest(request.digest, path_prefix=tmpdir_relative_path)
         result = interactive_runner.run(
-            InteractiveProcess(argv=(exe_path,), env=request.env, run_in_workspace=True)
+            InteractiveProcess(argv=request.args, env=request.env, run_in_workspace=True)
         )
-
     return Repl(result.exit_code)