Add ability to run any PythonSourceField (pantsbuild#15849)

Co-authored-by: Eric Arellano <14852634+Eric-Arellano@users.noreply.github.com>

[ci skip-build-wheels]

docs/markdown/Python/python-goals/python-run-goal.md

@@ -6,15 +6,22 @@ hidden: false
 createdAt: "2020-03-16T16:19:56.403Z"
 updatedAt: "2022-01-29T16:45:29.511Z"
 ---
-To run an executable/script, use `./pants run` on a [`pex_binary`](doc:reference-pex_binary) target. (See [package](doc:python-package-goal) for more on the `pex_binary` target.)
+To run an executable/script, use `./pants run` on one of the following target types:
+
+* [`pex_binary`](doc:reference-pex_binary)
+* [`python_source`](doc:reference-python_source)
+
+(See [package](doc:python-package-goal) for more on the `pex_binary` target.)
 
 ```bash
+# A python_source target (usually referred to by the filename)
 $ ./pants run project/app.py
 ```
 
 or
 
 ```bash
+# A pex_binary target (must be referred to by target name)
 $ ./pants run project:app
 ```
 
@@ -36,6 +43,26 @@ The program will have access to the same environment used by the parent `./pants
 >
 > Run `./pants dependencies --transitive path/to/binary.py` to ensure that all the files you need are showing up, including for any [assets](doc:assets) you intend to use.
 
+Execution Semantics
+-------------------
+
+Running a `pex_binary` is equivalent to `package`-ing the target followed by executing the built PEX
+from the repo root.
+
+Running a `python_source` with the `run_goal_use_sandbox` field set to `True` (the default) runs your
+code in an ephemeral sandbox (temporary directory) with your firstparty code and and
+Pants-generated files (such as a `relocated_files` or `archive`) copied inside. If you are using
+generated files like this, you may need to set the `run_goal_use_sandbox` to `True` for file loading
+to work properly.
+
+Running a `python_source` with the `run_goal_use_sandbox` field set to `False` is equivalent to
+running the source directly (a la `python ...`) with the set of third-party dependencies exposed to
+the interpreter. This is comparable to using a virtual environment or Poetry to run your script
+(E.g. `venv/bin/python ...` or `poetry run python ...`). When scripts write in-repo files—such as
+Django's `manage.py makemigrations` - it is often necessary to set `run_goal_use_sandbox` to `False`
+so that the file is written into the expected location.
+
+
 Watching the filesystem
 -----------------------
 

src/python/pants/backend/python/goals/BUILD

@@ -25,7 +25,8 @@ python_tests(
             "tags": ["platform_specific_behavior"],
             "timeout": 300,
         },
-        "run_pex_binary_integration_test.py": {"timeout": 180},
+        "run_pex_binary_integration_test.py": {"timeout": 400},
+        "run_python_source_integration_test.py": {"timeout": 180},
         "setup_py_integration_test.py": {
             "dependencies": ["testprojects/src/python:native_directory"],
             "tags": ["platform_specific_behavior"],

src/python/pants/backend/python/goals/package_pex_binary.py

@@ -25,7 +25,6 @@
     PexStripEnvField,
     ResolvedPexEntryPoint,
     ResolvePexEntryPointRequest,
-    RunInSandboxField,
 )
 from pants.backend.python.util_rules.pex import CompletePlatforms, Pex, PexPlatforms
 from pants.backend.python.util_rules.pex_from_targets import PexFromTargetsRequest
@@ -71,7 +70,6 @@ class PexBinaryFieldSet(PackageFieldSet, RunFieldSet):
     execution_mode: PexExecutionModeField
     include_requirements: PexIncludeRequirementsField
     include_tools: PexIncludeToolsField
-    run_in_sandbox: RunInSandboxField
 
     @property
     def _execution_mode(self) -> PexExecutionMode:

src/python/pants/backend/python/goals/run_helper.py

@@ -0,0 +1,172 @@
+# Copyright 2022 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+from __future__ import annotations
+
+import os
+from typing import Iterable, Optional
+
+from pants.backend.python.subsystems.debugpy import DebugPy
+from pants.backend.python.target_types import (
+    ConsoleScript,
+    PexEntryPointField,
+    ResolvedPexEntryPoint,
+    ResolvePexEntryPointRequest,
+)
+from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
+from pants.backend.python.util_rules.local_dists import LocalDistsPex, LocalDistsPexRequest
+from pants.backend.python.util_rules.pex import Pex, PexRequest
+from pants.backend.python.util_rules.pex_environment import PexEnvironment
+from pants.backend.python.util_rules.pex_from_targets import (
+    InterpreterConstraintsRequest,
+    PexFromTargetsRequest,
+)
+from pants.backend.python.util_rules.python_sources import (
+    PythonSourceFiles,
+    PythonSourceFilesRequest,
+)
+from pants.core.goals.run import RunDebugAdapterRequest, RunRequest
+from pants.core.subsystems.debug_adapter import DebugAdapterSubsystem
+from pants.engine.addresses import Address
+from pants.engine.fs import Digest, MergeDigests
+from pants.engine.rules import Get, MultiGet, rule_helper
+from pants.engine.target import TransitiveTargets, TransitiveTargetsRequest
+
+
+def _in_chroot(relpath: str) -> str:
+    return os.path.join("{chroot}", relpath)
+
+
+@rule_helper
+async def _create_python_source_run_request(
+    address: Address,
+    *,
+    entry_point_field: PexEntryPointField,
+    pex_env: PexEnvironment,
+    run_in_sandbox: bool,
+    console_script: Optional[ConsoleScript] = None,
+    additional_pex_args: Iterable[str] = (),
+) -> RunRequest:
+    addresses = [address]
+    entry_point, transitive_targets = await MultiGet(
+        Get(
+            ResolvedPexEntryPoint,
+            ResolvePexEntryPointRequest(entry_point_field),
+        ),
+        Get(TransitiveTargets, TransitiveTargetsRequest(addresses)),
+    )
+
+    interpreter_constraints = await Get(
+        InterpreterConstraints, InterpreterConstraintsRequest(addresses)
+    )
+
+    pex_filename = (
+        address.generated_name.replace(".", "_") if address.generated_name else address.target_name
+    )
+    pex_get = Get(
+        Pex,
+        PexFromTargetsRequest(
+            addresses,
+            output_filename=f"{pex_filename}.pex",
+            internal_only=True,
+            include_source_files=False,
+            # `PEX_EXTRA_SYS_PATH` should contain this entry_point's module.
+            main=console_script or entry_point.val,
+            additional_args=(
+                *additional_pex_args,
+                # N.B.: Since we cobble together the runtime environment via PEX_EXTRA_SYS_PATH
+                # below, it's important for any app that re-executes itself that these environment
+                # variables are not stripped.
+                "--no-strip-pex-env",
+            ),
+        ),
+    )
+    sources_get = Get(
+        PythonSourceFiles, PythonSourceFilesRequest(transitive_targets.closure, include_files=True)
+    )
+    pex, sources = await MultiGet(pex_get, sources_get)
+
+    local_dists = await Get(
+        LocalDistsPex,
+        LocalDistsPexRequest(
+            addresses,
+            internal_only=True,
+            interpreter_constraints=interpreter_constraints,
+            sources=sources,
+        ),
+    )
+
+    input_digests = [
+        pex.digest,
+        local_dists.pex.digest,
+        # Note regarding not-in-sandbox mode: You might think that the sources don't need to be copied
+        # into the chroot when using inline sources. But they do, because some of them might be
+        # codegenned, and those won't exist in the inline source tree. Rather than incurring the
+        # complexity of figuring out here which sources were codegenned, we copy everything.
+        # The inline source roots precede the chrooted ones in PEX_EXTRA_SYS_PATH, so the inline
+        # sources will take precedence and their copies in the chroot will be ignored.
+        local_dists.remaining_sources.source_files.snapshot.digest,
+    ]
+    merged_digest = await Get(Digest, MergeDigests(input_digests))
+
+    complete_pex_env = pex_env.in_workspace()
+    args = complete_pex_env.create_argv(_in_chroot(pex.name), python=pex.python)
+
+    chrooted_source_roots = [_in_chroot(sr) for sr in sources.source_roots]
+    # The order here is important: we want the in-repo sources to take precedence over their
+    # copies in the sandbox (see above for why those copies exist even in non-sandboxed mode).
+    source_roots = [
+        *([] if run_in_sandbox else sources.source_roots),
+        *chrooted_source_roots,
+    ]
+    extra_env = {
+        **pex_env.in_workspace().environment_dict(python_configured=pex.python is not None),
+        "PEX_PATH": _in_chroot(local_dists.pex.name),
+        "PEX_EXTRA_SYS_PATH": os.pathsep.join(source_roots),
+    }
+
+    return RunRequest(
+        digest=merged_digest,
+        args=args,
+        extra_env=extra_env,
+    )
+
+
+@rule_helper
+async def _create_python_source_run_dap_request(
+    regular_run_request: RunRequest,
+    *,
+    entry_point_field: PexEntryPointField,
+    debugpy: DebugPy,
+    debug_adapter: DebugAdapterSubsystem,
+    console_script: Optional[ConsoleScript] = None,
+) -> RunDebugAdapterRequest:
+    entry_point, debugpy_pex = await MultiGet(
+        Get(
+            ResolvedPexEntryPoint,
+            ResolvePexEntryPointRequest(entry_point_field),
+        ),
+        Get(Pex, PexRequest, debugpy.to_pex_request()),
+    )
+
+    merged_digest = await Get(
+        Digest, MergeDigests([regular_run_request.digest, debugpy_pex.digest])
+    )
+    extra_env = dict(regular_run_request.extra_env)
+    extra_env["PEX_PATH"] = os.pathsep.join(
+        [
+            extra_env["PEX_PATH"],
+            # For debugpy to work properly, we need to have just one "environment" for our
+            # command to run in. Therefore, we cobble one together by exeucting debugpy's PEX, and
+            # shoehorning in the original PEX through PEX_PATH.
+            _in_chroot(os.path.basename(regular_run_request.args[1])),
+        ]
+    )
+    main = console_script or entry_point.val
+    assert main is not None
+    args = [
+        regular_run_request.args[0],  # python executable
+        _in_chroot(debugpy_pex.name),
+        *debugpy.get_args(debug_adapter, main),
+    ]
+
+    return RunDebugAdapterRequest(digest=merged_digest, args=args, extra_env=extra_env)