Use appropriate shebang for multi-platform PEXes. (pex-tool#2296)

Although it's not always possible to derive an appropriate shebang that
will work for multi-platform PEXes, we now do so when possible and warn
when we cannot.

Fixes pex-tool#1540

pex/bin/pex.py

@@ -872,8 +872,24 @@ def build_pex(
             filename=options.executable, env_filename="__pex_executable__.py"
         )
 
-    if options.python_shebang:
-        pex_builder.set_shebang(options.python_shebang)
+    specific_shebang = options.python_shebang or targets.compatible_shebang()
+    if specific_shebang:
+        pex_builder.set_shebang(specific_shebang)
+    else:
+        # TODO(John Sirois): Consider changing fallback to `#!/usr/bin/env python` in Pex 3.x.
+        pex_warnings.warn(
+            "Could not calculate a targeted shebang for:\n"
+            "{targets}\n"
+            "\n"
+            "Using shebang: {default_shebang}\n"
+            "If this is not appropriate, you can specify a custom shebang using the "
+            "--python-shebang option.".format(
+                targets="\n".join(
+                    sorted(target.render_description() for target in targets.unique_targets())
+                ),
+                default_shebang=pex_builder.shebang,
+            )
+        )
 
     return pex_builder
 

pex/pex_builder.py

@@ -457,6 +457,11 @@ def set_entry_point(self, entry_point):
         self._ensure_unfrozen("Setting an entry point")
         self._pex_info.entry_point = entry_point
 
+    @property
+    def shebang(self):
+        # type: () -> str
+        return self._shebang
+
     def set_shebang(self, shebang):
         """Set the exact shebang line for the PEX file.
 

pex/targets.py

@@ -5,6 +5,7 @@
 
 import os
 
+from pex import pex_warnings
 from pex.dist_metadata import Requirement
 from pex.interpreter import PythonInterpreter, calculate_binary_name
 from pex.orderedset import OrderedSet
@@ -377,3 +378,16 @@ def require_at_most_one_target(self, purpose):
             return cast(Target, next(iter(resolved_targets)))
         except StopIteration:
             return None
+
+    def compatible_shebang(self):
+        # type: () -> Optional[str]
+        pythons = {
+            (target.platform.impl, target.platform.version_info[:2])
+            for target in self.unique_targets()
+        }
+        if len(pythons) == 1:
+            impl, version = pythons.pop()
+            return "#!/usr/bin/env {python}{version}".format(
+                python="pypy" if impl == "pp" else "python", version=".".join(map(str, version))
+            )
+        return None

tests/integration/test_issue_1540.py

@@ -0,0 +1,85 @@
+# Copyright 2023 Pants project contributors (see CONTRIBUTORS.md).
+# Licensed under the Apache License, Version 2.0 (see LICENSE).
+
+from __future__ import absolute_import
+
+import os.path
+from textwrap import dedent
+
+from pex.interpreter import PythonInterpreter
+from pex.typing import TYPE_CHECKING
+from testing import run_pex_command
+
+if TYPE_CHECKING:
+    from typing import Any
+
+
+def test_derive_consistent_shebang_platforms(
+    tmpdir,  # type: Any
+    current_interpreter,  # type: PythonInterpreter
+):
+    # type: (...) -> None
+
+    pex = os.path.join(str(tmpdir), "pex")
+
+    def read_pex_shebang():
+        # type: () -> bytes
+        with open(pex, "rb") as fp:
+            return fp.readline()
+
+    run_pex_command(args=["--platform", "linux_x86_64-cp-311-cp311", "-o", pex]).assert_success()
+    assert b"#!/usr/bin/env python3.11\n" == read_pex_shebang()
+
+    run_pex_command(
+        args=[
+            "--platform",
+            "linux_x86_64-cp-311-cp311",
+            "--platform",
+            "macosx_10.9_x86_64-cp-311-cp311",
+            "-o",
+            pex,
+        ]
+    ).assert_success()
+    assert b"#!/usr/bin/env python3.11\n" == read_pex_shebang()
+
+    run_pex_command(
+        args=[
+            "--platform",
+            "linux_x86_64-cp-3.11.5-cp311",
+            "--platform",
+            "macosx_10.9_x86_64-cp-311-cp311",
+            "-o",
+            pex,
+        ]
+    ).assert_success()
+    assert b"#!/usr/bin/env python3.11\n" == read_pex_shebang()
+
+    result = run_pex_command(
+        args=[
+            "--platform",
+            "linux_x86_64-cp-310-cp310",
+            "--platform",
+            "macosx_10.9_x86_64-cp-311-cp311",
+            "-o",
+            pex,
+        ]
+    )
+    result.assert_success()
+    current_interpreter_shebang = current_interpreter.identity.hashbang()
+    assert (
+        "{shebang}\n".format(shebang=current_interpreter_shebang).encode("utf-8")
+        == read_pex_shebang()
+    )
+    assert (
+        dedent(
+            """\
+            PEXWarning: Could not calculate a targeted shebang for:
+            abbreviated platform cp310-cp310-linux_x86_64
+            abbreviated platform cp311-cp311-macosx_10_9_x86_64
+
+            Using shebang: {shebang}
+            If this is not appropriate, you can specify a custom shebang using the --python-shebang option.
+            """
+        ).format(shebang=current_interpreter_shebang)
+        in result.error
+    )

tests/test_targets.py

@@ -24,6 +24,7 @@
 )
 from pex.third_party.packaging.specifiers import SpecifierSet
 from pex.typing import TYPE_CHECKING
+from testing import IS_PYPY, PY39, PY310, PY_VER, ensure_python_interpreter
 
 if TYPE_CHECKING:
     from typing import Optional
@@ -261,3 +262,70 @@ def test_from_target_complete_platform(current_interpreter):
     assert tgts.interpreter is None
     assert OrderedSet([complete_platform]) == tgts.unique_targets(only_explicit=False)
     assert OrderedSet([complete_platform]) == tgts.unique_targets(only_explicit=True)
+
+
+def test_compatible_shebang(
+    current_interpreter,  # type: PythonInterpreter
+    current_platform,  # type: Platform
+):
+    # type: (...) -> None
+
+    current_python = "pypy" if IS_PYPY else "python"
+    current_version = ".".join(map(str, PY_VER))
+    current_python_shebang = "#!/usr/bin/env {python}{version}".format(
+        python=current_python, version=current_version
+    )
+    assert (
+        current_python_shebang
+        == Targets.from_target(LocalInterpreter.create(current_interpreter)).compatible_shebang()
+    )
+    assert (
+        current_python_shebang
+        == Targets.from_target(AbbreviatedPlatform.create(current_platform)).compatible_shebang()
+    )
+    current_complete_platform = CompletePlatform.from_interpreter(current_interpreter)
+    assert (
+        current_python_shebang
+        == Targets.from_target(current_complete_platform).compatible_shebang()
+    )
+    assert (
+        current_python_shebang
+        == Targets(
+            interpreters=(current_interpreter,),
+            complete_platforms=(current_complete_platform,),
+            platforms=(current_platform,),
+        ).compatible_shebang()
+    )
+
+    incompatible_interpreter = PythonInterpreter.from_binary(
+        ensure_python_interpreter(PY39 if PY_VER == (3, 10) else PY310)
+    )
+    assert (
+        Targets(interpreters=(current_interpreter, incompatible_interpreter)).compatible_shebang()
+        is None
+    )
+
+    incompatible_version_info = (3, 9) if PY_VER == (3, 10) else (3, 10)
+    incompatible_platform_version = Platform(
+        platform=current_platform.platform,
+        impl=current_platform.impl,
+        version_info=incompatible_version_info,
+        version=".".join(map(str, incompatible_version_info)),
+        abi=current_platform.abi,
+    )
+    assert (
+        Targets(platforms=(current_platform, incompatible_platform_version)).compatible_shebang()
+        is None
+    )
+
+    incompatible_platform_impl = Platform(
+        platform=current_platform.platform,
+        impl="cp" if current_platform.impl == "pp" else "pp",
+        version_info=current_platform.version_info,
+        version=current_platform.version,
+        abi=current_platform.abi,
+    )
+    assert (
+        Targets(platforms=(current_platform, incompatible_platform_impl)).compatible_shebang()
+        is None
+    )