Keep testrunner dependencies separate and set shebang in config

.plzconfig

@@ -93,6 +93,16 @@ DefaultValue = //tools:wheel_resolver
 Optional = true
 Inherit = true
 
+[PluginConfig "pex_shebang"]
+ConfigKey = PexShebang
+DefaultValue = "#!/usr/bin/python3"
+Optional = true
+Inherit = true
+
+[PluginConfig "portalocker"]
+ConfigKey = Portalocker
+DefaultValue = //third_party/python:py_binary_bootstrap
+
 [featureflags]
 PythonWheelHashing = true
 ExcludePythonRules = true

build_defs/python.build_defs

@@ -131,7 +131,7 @@ def python_binary(name:str, main:str, srcs:list=[], resources:list=[], out:str=N
       labels (list): Labels to apply to this rule.
     """
     interpreter = interpreter or CONFIG.PYTHON.DEFAULT_INTERPRETER
-    shebang = shebang or _set_shebang(interpreter, CONFIG.PYTHON.INTERPRETER_OPTIONS, site)
+    shebang = shebang or CONFIG.PYTHON.PEX_SHEBANG
     zipsafe_flag = '' if zip_safe is False else '--zip_safe'
     cmd = '$TOOLS_PEX -m "%s" %s --interpreter_options="%s" --stamp="$RULE_HASH"' % (
         CONFIG.PYTHON.MODULE_DIR, zipsafe_flag, CONFIG.PYTHON.INTERPRETER_OPTIONS)
@@ -175,15 +175,21 @@ def python_binary(name:str, main:str, srcs:list=[], resources:list=[], out:str=N
     )
 
     # This rule concatenates the .pex with all the other precompiled zip files from dependent rules.
-    cmd = f'mv $SRC __main__.py && $TOOL z -i . -s .pex.zip -s .whl --preamble="{shebang}" --include_other --add_init_py --strict'
+    cmd = f'mv $SRCS_MAIN __main__.py && mv $SRCS_BOOTSTRAP . && $TOOL z -i . -s .zip -s .whl --preamble="{shebang}" --include_other --add_init_py --strict'
+
     if strip:
         cmd += ' --strip_py'
 
     debug_cmd = _debug_cmd("./$OUT")
 
+    srcs = {
+        'main': [pex_rule],
+        'bootstrap': [CONFIG.PYTHON.PORTALOCKER],
+        }
+
     return build_rule(
         name=name,
-        srcs=[pex_rule],
+        srcs=srcs,
         deps=[lib_rule],
         outs=[out or (name + '.pex')],
         data=data,
@@ -220,7 +226,7 @@ def _set_shebang(interpreter:str, interpreter_options:str, site:bool):
 
 def python_test(name:str, srcs:list, data:list|dict=[], resources:list=[], deps:list=[], worker:str='',
                 labels:list&features&tags=[], size:str=None, flags:str='', visibility:list=None,
-                sandbox:bool=None, timeout:int=0, flaky:bool|int=0,
+                sandbox:bool=None, timeout:int=0, flaky:bool|int=0, shebang:str=None,
                 test_outputs:list=None, zip_safe:bool=None, interpreter:str=None, site:bool=False,
                 test_runner:str=None):
     """Generates a Python test target.
@@ -288,7 +294,7 @@ def python_test(name:str, srcs:list, data:list|dict=[], resources:list=[], deps:
         pre_build=_handle_zip_safe,
         deps=deps,
         tools={
-            'interpreter': [interpreter or CONFIG.PYTHON.DEFAULT_INTERPRETER],
+            'interpreter': [interpreter],
             'pex': [CONFIG.PYTHON.PEX_TOOL],
         },
         labels = labels,
@@ -307,13 +313,16 @@ def python_test(name:str, srcs:list, data:list|dict=[], resources:list=[], deps:
     )
 
     deps = [pex_rule, lib_rule]
-    shebang = _set_shebang(interpreter, interpreter_options, site)
+    shebang = shebang or CONFIG.PYTHON.PEX_SHEBANG
+
+    if not site:
+        interpreter_options = f'{interpreter_options} -S'
 
     # If there are resources specified, they have to get built into the pex.
     # Also add the the value of CONFIG.PYTHON.TESTRUNNER_BOOTSTRAP as a dependency.
     deps += [CONFIG.PYTHON.TESTRUNNER_DEPS]
 
-    test_cmd = f'$TEST {flags}'
+    test_cmd = f'$TOOL {interpreter_options} $TEST {flags}'
     worker_cmd = f'$(worker {worker})' if worker else ""
     if worker_cmd:
         test_cmd = f'{worker_cmd} && {test_cmd} '
@@ -325,17 +334,25 @@ def python_test(name:str, srcs:list, data:list|dict=[], resources:list=[], deps:
         pre_cmd = worker_cmd,
     )
 
+    #'cp -r $SRCS_TESTRUNNERDEPS .',
+    cmd = ['mv $SRCS_MAIN __main__.py',
+           'mv $SRCS_TESTRUNNERDEPS .',
+           f'$TOOL z -i . -s .pex.zip -s .zip -s .whl --preamble="{shebang}" --include_other --add_init_py --strict']
+
     # This rule concatenates the .pex with all the other precompiled zip files from dependent rules.
     return build_rule(
         name=name,
-        srcs=[pex_rule],
+        srcs={
+            "main": pex_rule,
+            "testrunnerdeps" : CONFIG.PYTHON.TESTRUNNER_DEPS
+        },
         deps=deps,
         # N.B. the actual test sources are passed as data files as well. This is needed for pytest but
         #      is faster for unittest as well (because we don't need to rebuild the pex if they change).
         data=data | {'_srcs': srcs} if isinstance(data, dict) else data + srcs,
         outs=[f'{name}.pex'],
         labels=labels + ['test_results_dir'],
-        cmd=f'mv $SRC __main__.py && $TOOL z -i . -s .pex.zip -s .whl --preamble="{shebang}" --include_other --add_init_py --strict',
+        cmd=' && '.join(cmd),
         test_cmd=test_cmd,
         debug_cmd=debug_cmd,
         needs_transitive_deps=True,
@@ -351,6 +368,7 @@ def python_test(name:str, srcs:list, data:list|dict=[], resources:list=[], deps:
         test_outputs=test_outputs,
         requires=['py', 'test', interpreter or CONFIG.PYTHON.DEFAULT_INTERPRETER],
         tools=[CONFIG.JARCAT_TOOL],
+        test_tools=[CONFIG.PYTHON.DEFAULT_INTERPRETER],
     )
 
 def pip_library(name:str, version:str, hashes:list=None, package_name:str=None,

third_party/python/BUILD

@@ -176,8 +176,14 @@ pip_library(
     deps = [":importlib_metadata"],
 )
 
-pip_library(
+python_wheel(
     name = "importlib_metadata",
+    outs = [
+        "importlib_metadata",
+        "importlib_metadata-1.5.0.dist-info",
+    ],
+    hashes = ["b97607a1a18a5100839aec1dc26a1ea17ee0d93b20b0f008d80a5a050afb200b"],
+    patch = "importlib_metadata.patch",
     version = "1.5.0",
     deps = [":zipp"],
 )
@@ -299,6 +305,24 @@ pip_library(
     version = "1.7.0",
 )
 
+genrule(
+    name = "py_binary_bootstrap",
+    cmd = [
+        "mkdir .bootstrap",
+        "ls -d $PKG_DIR/* | grep -v \.whl | xargs -I{} mv {} .bootstrap;"
+        "ls -d $PKG_DIR/* | xargs -I{} -n1 $TOOLS x {} -o {}.unzip;"
+        "for dir in $PKG_DIR/*.unzip; do mv $dir/$PKG_DIR/* .bootstrap; done",
+        "$TOOLS z -i .bootstrap/ -o .py_binary_bootstrap.zip  --include_other",
+    ],
+    srcs = [
+        ":portalocker"
+    ],
+    outs = [
+        ".py_binary_bootstrap.zip",
+    ],
+    tools = [CONFIG.JARCAT_TOOL],
+)
+
 pip_library(
     name = "numpy",
     test_only = True,
@@ -469,8 +493,26 @@ pip_library(
     version = "1.5.0",
 )
 
-filegroup(
+genrule(
     name = "unittest_bootstrap",
+    cmd = [
+        "mkdir .bootstrap",
+        "ls -d $PKG_DIR/* | grep -v \.whl | xargs -I{} mv {} .bootstrap;"
+        "ls -d $PKG_DIR/* | xargs -I{} -n1 $TOOLS x {} -o {}.unzip;"
+        "for dir in $PKG_DIR/*.unzip; do mv $dir/$PKG_DIR/* .bootstrap; done",
+        "$TOOLS z -i .bootstrap/ -o .unittest_bootstrap.zip  --include_other",
+    ],
+    srcs = [
+        ":unittest_deps"
+    ],
+    outs = [
+        ".unittest_bootstrap.zip",
+    ],
+    tools = [CONFIG.JARCAT_TOOL],
+)
+
+filegroup(
+    name = "unittest_deps",
     srcs = [
         ":coverage",
         ":portalocker",
@@ -479,8 +521,26 @@ filegroup(
     ],
 )
 
-filegroup(
+genrule(
     name = "pytest_bootstrap",
+    cmd = [
+        "mkdir .bootstrap",
+        "ls -d $PKG_DIR/* | grep -v \.whl | xargs -I{} mv {} .bootstrap;"
+        "ls -d $PKG_DIR/* | xargs -I{} -n1 $TOOLS x {} -o {}.unzip;"
+        "for dir in $PKG_DIR/*.unzip; do mv $dir/$PKG_DIR/* .bootstrap; done",
+        "$TOOLS z -i .bootstrap/ -o .pytest_bootstrap.zip  --include_other",
+    ],
+    srcs = [
+        ":pytest_deps"
+    ],
+    outs = [
+        ".pytest_bootstrap.zip",
+    ],
+    tools = [CONFIG.JARCAT_TOOL],
+)
+
+filegroup(
+    name = "pytest_deps",
     srcs = [
         ":attrs",
         ":funcsigs",

tools/please_pex/pex/pex_main.py

@@ -113,7 +113,6 @@ class SoImport(object):
     """So import. Much binary. Such dynamic. Wow."""
 
     def __init__(self):
-
         if PY_VERSION.major < 3:
             self.suffixes = {x[0]: x for x in imp.get_suffixes() if x[2] == imp.C_EXTENSION}
         else:
@@ -127,6 +126,8 @@ def __init__(self):
             for name in zf.namelist():
                 path, _ = self.splitext(name)
                 if path:
+                    if path.startswith('.bootstrap/'):
+                        path = path[len('.bootstrap/'):]
                     importpath = path.replace('/', '.')
                     self.modules.setdefault(importpath, name)
                     if path.startswith(MODULE_DIR):
@@ -281,8 +282,8 @@ def explode_zip():
     inside a zipfile.
     """
     # Temporarily add bootstrap to sys path
-    ### THIS SHOULD MAYBE COME FROM CONFIG FILE?
-    sys.path = [os.path.join(sys.path[0], 'third_party/python')] + sys.path[1:]
+    sys.path = [os.path.join(sys.path[0], '.bootstrap')] + sys.path[1:]
+    print("sys.path =", sys.path)
     import contextlib, portalocker
     sys.path = sys.path[1:]
 
@@ -360,6 +361,8 @@ def main():
 
     N.B. This gets redefined by pex_test_main to run tests instead.
     """
+    # Add .bootstrap dir to path, after the initial pex entry
+    sys.path = sys.path[:1] + [os.path.join(sys.path[0], '.bootstrap')] + sys.path[1:]
     # Starts a debugging session, if defined, before running the entry point.
     start_debugger()
     # Must run this as __main__ so it executes its own __name__ == '__main__' block.

tools/please_pex/pex/pex_run.py

@@ -1,5 +1,5 @@
-def run():
-    if not ZIP_SAFE:
+def run(explode=False):
+    if explode or not ZIP_SAFE:
         with explode_zip()():
             add_module_dir_to_sys_path(MODULE_DIR)
             return main()
@@ -10,15 +10,18 @@ def run():
 
 
 if __name__ == '__main__':
+     # If PEX_EXPLODE is set, then it should always be exploded.
+    explode = os.environ.get('PEX_EXPLODE', '0') != '0'
+
     # If PEX_INTERPRETER is set, then it starts an interactive console.
     if os.environ.get('PEX_INTERPRETER', '0') != '0':
         import code
         result = code.interact()
     # If PEX_PROFILE_FILENAME is set, then it collects profile information into the filename.
     elif os.environ.get('PEX_PROFILE_FILENAME'):
         with profile(os.environ['PEX_PROFILE_FILENAME'])():
-            result = run()
+            result = run(explode)
     else:
-        result = run()
+        result = run(explode)
 
     sys.exit(result)

tools/please_pex/pex/pex_test_main.py

@@ -31,6 +31,7 @@ def _xml_file(self, fr, analysis, *args, **kvargs):
 def main():
     """Runs the tests. Returns an appropriate exit code."""
     args = [arg for arg in sys.argv[1:]]
+    sys.path = sys.path[:1] + [os.path.join(sys.path[0], '.bootstrap')] + sys.path[1:]
     if os.getenv('COVERAGE'):
         # It's important that we run coverage while we load the tests otherwise
         # we get no coverage for import statements etc.