Add ability to run any PythonSourceField

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

@@ -6,15 +6,21 @@ 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 any 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 +42,21 @@ 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` 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. `ve/bin/python ...` or `poetry run python ...`).
+
+If your `python_source` depends on Pants-generated files (such as a `relocated_files` or `archive`)
+you might consider setting the `run_goal_use_sandbox` field on your `python_source` to `True`.
+Pants only materializes those files in the sandbox and not in your repo, so your code might have
+trouble locating them.
+
 Watching the filesystem
 -----------------------
 

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,127 @@
+# Copyright 2022 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+import os
+from typing import Iterable, Optional
+
+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
+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 RunRequest
+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
+
+
+@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))
+
+    def in_chroot(relpath: str) -> str:
+        return os.path.join("{chroot}", relpath)
+
+    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,
+    )

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

@@ -4,122 +4,41 @@
 import os
 
 from pants.backend.python.goals.package_pex_binary import PexBinaryFieldSet
-from pants.backend.python.target_types import (
-    PexBinaryDefaults,
-    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
-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.backend.python.goals.run_helper import _create_python_source_run_request
+from pants.backend.python.target_types import PexBinaryDefaults
+from pants.backend.python.util_rules.pex_environment import PexEnvironment, PexRuntimeEnvironment
+from pants.core.goals.package import BuiltPackage
 from pants.core.goals.run import RunFieldSet, RunRequest
-from pants.engine.fs import Digest, MergeDigests
-from pants.engine.rules import Get, MultiGet, collect_rules, rule
-from pants.engine.target import TransitiveTargets, TransitiveTargetsRequest
+from pants.engine.rules import Get, collect_rules, rule
 from pants.engine.unions import UnionRule
 from pants.util.logging import LogLevel
 
 
 @rule(level=LogLevel.DEBUG)
 async def create_pex_binary_run_request(
-    field_set: PexBinaryFieldSet, pex_binary_defaults: PexBinaryDefaults, pex_env: PexEnvironment
+    field_set: PexBinaryFieldSet,
+    pex: PexRuntimeEnvironment,
+    pex_binary_defaults: PexBinaryDefaults,
+    pex_env: PexEnvironment,
 ) -> RunRequest:
-    run_in_sandbox = field_set.run_in_sandbox.value
-    entry_point, transitive_targets = await MultiGet(
-        Get(
-            ResolvedPexEntryPoint,
-            ResolvePexEntryPointRequest(field_set.entry_point),
-        ),
-        Get(TransitiveTargets, TransitiveTargetsRequest([field_set.address])),
+    if pex.run_packaged_firstparty_code:
+        built_pex = await Get(BuiltPackage, PexBinaryFieldSet, field_set)
+        relpath = built_pex.artifacts[0].relpath
+        assert relpath is not None
+        return RunRequest(
+            digest=built_pex.digest,
+            args=[os.path.join("{chroot}", relpath)],
+        )
+
+    return await _create_python_source_run_request(
+        field_set.address,
+        entry_point_field=field_set.entry_point,
+        pex_env=pex_env,
+        run_in_sandbox=True,
+        console_script=field_set.script.value,
+        additional_pex_args=field_set.generate_additional_args(pex_binary_defaults),
     )
 
-    addresses = [field_set.address]
-    interpreter_constraints = await Get(
-        InterpreterConstraints, InterpreterConstraintsRequest(addresses)
-    )
-
-    pex_filename = (
-        field_set.address.generated_name.replace(".", "_")
-        if field_set.address.generated_name
-        else field_set.address.target_name
-    )
-    pex_get = Get(
-        Pex,
-        PexFromTargetsRequest(
-            [field_set.address],
-            output_filename=f"{pex_filename}.pex",
-            internal_only=True,
-            include_source_files=False,
-            # Note that the file for first-party entry points is not in the PEX itself. In that
-            # case, it's loaded by setting `PEX_EXTRA_SYS_PATH`.
-            main=entry_point.val or field_set.script.value,
-            additional_args=(
-                *field_set.generate_additional_args(pex_binary_defaults),
-                # 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(
-            [field_set.address],
-            internal_only=True,
-            interpreter_constraints=interpreter_constraints,
-            sources=sources,
-        ),
-    )
-
-    input_digests = [
-        pex.digest,
-        local_dists.pex.digest,
-        # Note regarding inline 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))
-
-    def in_chroot(relpath: str) -> str:
-        return os.path.join("{chroot}", relpath)
-
-    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 = {
-        **complete_pex_env.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)
-
 
 def rules():
     return [*collect_rules(), UnionRule(RunFieldSet, PexBinaryFieldSet)]