Repository: pex-tool/pex

Issue:
Issue when running a module inside pex file
I met some strange behaviors when running a module inside pex file. I found a simple reproducible example as below.

Python: 3.8.3
Pex: 2.1.15

I have a `test.py` under a folder `src`,
```
[cwd]
└── src
    └── test.py
```

`test.py`:
```python
import pickle

def add(a: int, b: int) -> int:
    return a+b

def main() -> None:

    pickle_add = pickle.dumps(add)
    add_clone = pickle.loads(pickle_add)
    print(add_clone(1, 2))

if __name__ == "__main__":
    main()
```

To build a pex file with entry point:
    
    pex -D ./src -e test:main -o test_with_ep.pex

To build a pex file without entry point:

    pex -D ./src -o test_no_ep.pex

No issue when running:

    python test_with_ep.pex      
    3

But when running `test_no_ep.pex`:

    python test_no_ep.pex -m test
    Traceback (most recent call last):
      File "/private/tmp/test_no_ep.pex/.bootstrap/pex/pex.py", line 423, in execute
      File "/private/tmp/test_no_ep.pex/.bootstrap/pex/pex.py", line 356, in _wrap_coverage
      File "/private/tmp/test_no_ep.pex/.bootstrap/pex/pex.py", line 387, in _wrap_profiling
      File "/private/tmp/test_no_ep.pex/.bootstrap/pex/pex.py", line 477, in _execute
      File "/private/tmp/test_no_ep.pex/.bootstrap/pex/pex.py", line 520, in execute_interpreter
      File "/private/tmp/test_no_ep.pex/.bootstrap/pex/pex.py", line 584, in execute_module
      File "/Users/yuhan/.pyenv/versions/3.8.3/lib/python3.8/runpy.py", line 210, in run_module
        return _run_code(code, {}, init_globals, run_name, mod_spec)
      File "/Users/yuhan/.pyenv/versions/3.8.3/lib/python3.8/runpy.py", line 87, in _run_code
        exec(code, run_globals)
      File "/private/tmp/test_no_ep.pex/test.py", line 13, in <module>
      File "/private/tmp/test_no_ep.pex/test.py", line 8, in main
    _pickle.PicklingError: Can't pickle <function add at 0x1042aef70>: attribute lookup add on __main__ failed

Problem could be solved if I move that function `main()` into another file and import it in `test.py`. Then I could run that module in pex file.

Do I miss something here ?


Discussion:
Hmm, well, I can reproduce this at least, so it's consistent. Will take a look.
This looks like a quirk of how pickle works with module-level objects, possibly related to https://bugs.python.org/issue37852. 
It looks like the issue here is Pex use of `runpy.run_module`. Passing `alter_sys=True`, which is what `python -m` does, solves this.

See the CPython main -m handling here which uses `runpy.run_module`:
https://github.com/python/cpython/blob/ba251c2ae6654bfc8abd9d886b219698ad34ac3c/Modules/main.c#L580-L581
https://github.com/python/cpython/blob/ba251c2ae6654bfc8abd9d886b219698ad34ac3c/Modules/main.c#L258-L283
Presumably this would not affect pexes built with --venv?
> Presumably this would not affect pexes built with --venv?

Nope - it also affects --venv mode. Both normal and --venv still insert a `__main__` that delegates to the configured entry point and both delegations suffer the same issue.
Aha, that is kind of good news as it might explain some weirdness I'm seeing when running mypy on its own code from a --venv pex 

It's a bit of a stretch though.

Resolving patch:
diff --git a/pex/pex.py b/pex/pex.py
--- a/pex/pex.py
+++ b/pex/pex.py
@@ -7,6 +7,7 @@
 import re
 import sys
 import warnings
+import zipfile
 from distutils import sysconfig
 from site import USER_SITE
 from types import ModuleType
@@ -576,7 +577,7 @@ def log(msg, V=1):
         log("  * - paths that do not exist or will be imported via zipimport")
 
     def execute_interpreter(self):
-        # type: () -> Optional[str]
+        # type: () -> Any
         args = sys.argv[1:]
         if args:
             # NB: We take care here to setup sys.argv to match how CPython does it for each case.
@@ -588,7 +589,7 @@ def execute_interpreter(self):
             elif arg == "-m":
                 module = args[1]
                 sys.argv = args[1:]
-                return self.execute_module(module)
+                return self.execute_module(module, alter_sys=True)
             else:
                 try:
                     if arg == "-":
@@ -654,20 +655,33 @@ def execute_content(
         exec_function(ast, globals_map)
         return None
 
-    @classmethod
-    def execute_entry(cls, entry_point):
+    def execute_entry(self, entry_point):
         # type: (str) -> Any
-        runner = cls.execute_pkg_resources if ":" in entry_point else cls.execute_module
-        return runner(entry_point)
-
-    @classmethod
-    def execute_module(cls, module_name):
-        # type: (str) -> None
-        cls.demote_bootstrap()
+        if ":" in entry_point:
+            return self.execute_pkg_resources(entry_point)
+
+        # When running as a zipapp we can't usefully `alter_sys` to point to the module we're
+        # executing as argv[0] since that module will be contained in a zipfile. In other words, if
+        # we did do this, sys.argv[0] would be `/path/to/pex.zip/path/to/module.py` and that value
+        # would not be usable in the standard way. Its not a file you can read or re-execute
+        # against like a loose python module source file would be.
+        #
+        # Python itself disallows this case altogether in the standard zipapp module (probably for
+        # similar reasons). See: https://docs.python.org/3/library/zipapp.html#cmdoption-zipapp-m
+        alter_sys = not zipfile.is_zipfile(self._pex)
+        return self.execute_module(entry_point, alter_sys)
+
+    def execute_module(
+        self,
+        module_name,  # type: str
+        alter_sys,  # type: bool
+    ):
+        # type: (...) -> None
+        self.demote_bootstrap()
 
         import runpy
 
-        runpy.run_module(module_name, run_name="__main__")
+        runpy.run_module(module_name, run_name="__main__", alter_sys=alter_sys)
 
     @classmethod
     def execute_pkg_resources(cls, spec):
diff --git a/pex/tools/commands/venv.py b/pex/tools/commands/venv.py
--- a/pex/tools/commands/venv.py
+++ b/pex/tools/commands/venv.py
@@ -238,7 +238,7 @@ def record_provenance(src_to_dst):
             module_name, _, function = entry_point.partition(":")
             if not function:
                 import runpy
-                runpy.run_module(module_name, run_name="__main__")
+                runpy.run_module(module_name, run_name="__main__", alter_sys=True)
             else:
                 import importlib
                 module = importlib.import_module(module_name)


Test patch:
diff --git a/pex/testing.py b/pex/testing.py
--- a/pex/testing.py
+++ b/pex/testing.py
@@ -44,6 +44,7 @@
 
 PY_VER = sys.version_info[:2]
 IS_PYPY = hasattr(sys, "pypy_version_info")
+IS_PYPY2 = IS_PYPY and sys.version_info[0] == 2
 IS_PYPY3 = IS_PYPY and sys.version_info[0] == 3
 NOT_CPYTHON27 = IS_PYPY or PY_VER != (2, 7)
 NOT_CPYTHON36 = IS_PYPY or PY_VER != (3, 6)
diff --git a/tests/test_integration.py b/tests/test_integration.py
--- a/tests/test_integration.py
+++ b/tests/test_integration.py
@@ -42,6 +42,7 @@
 from pex.requirements import LogicalLine, PyPIRequirement, URLFetcher, parse_requirement_file
 from pex.testing import (
     IS_PYPY,
+    IS_PYPY2,
     NOT_CPYTHON27,
     NOT_CPYTHON27_OR_OSX,
     NOT_CPYTHON36_OR_LINUX,
@@ -2324,7 +2325,7 @@ def test_unzip_mode():
         with safe_open(os.path.join(src_dir, "example.py"), "w") as fp:
             fp.write(
                 dedent(
-                    """
+                    """\
                     import os
                     import sys
 
@@ -2334,7 +2335,7 @@ def test_unzip_mode():
 
                     print(' '.join(sys.argv[1:]))
                     sys.stdout.flush()
-                    os.execv(sys.argv[0], sys.argv[:-1])
+                    os.execv(sys.executable, [sys.executable] + sys.argv[:-1])
                     """
                 )
             )
@@ -2358,18 +2359,22 @@ def test_unzip_mode():
         output1 = subprocess.check_output(
             args=[pex_file, "quit", "re-exec"],
         )
-        assert ["quit re-exec", os.path.realpath(pex_file)] == output1.decode("utf-8").splitlines()
 
         pex_hash = PexInfo.from_pex(pex_file).pex_hash
         assert pex_hash is not None
         unzipped_cache = unzip_dir(pex_root, pex_hash)
         assert os.path.isdir(unzipped_cache)
+        assert [
+            "quit re-exec",
+            os.path.realpath(os.path.join(unzipped_cache, "example.py")),
+        ] == output1.decode("utf-8").splitlines()
+
         shutil.rmtree(unzipped_cache)
 
         output2 = subprocess.check_output(
             args=[pex_file, "quit", "re-exec"], env=make_env(PEX_UNZIP=False)
         )
-        assert output1 == output2
+        assert ["quit re-exec", os.path.realpath(pex_file)] == output2.decode("utf-8").splitlines()
         assert not os.path.exists(unzipped_cache)
 
 
@@ -3210,3 +3215,93 @@ def test_console_script_from_pex_path(tmpdir):
     result.assert_success()
 
     assert "hello world!\n" == subprocess.check_output(args=[pex_file]).decode("utf-8")
+
+
+def test_execute_module_issues_1018(tmpdir):
+    # type: (Any) -> None
+    src_dir = os.path.join(str(tmpdir), "src")
+    with safe_open(os.path.join(src_dir, "issues_1018.py"), "w") as fp:
+        fp.write(
+            dedent(
+                """\
+                import pickle
+
+                def add(a, b):
+                    return a + b
+
+                def main():
+                    pickle_add = pickle.dumps(add)
+                    add_clone = pickle.loads(pickle_add)
+                    print(add_clone(1, 2))
+
+                if __name__ == "__main__":
+                    main()
+                """
+            )
+        )
+    expected_output = b"3\n"
+
+    # There are 9 ways we can invoke the module above using Pex corresponding to the 2D matrix with
+    # axes: {zipapp, unzip, venv} x {entrypoint-function, entrypoint-module, ad-hoc (-m) module}.
+    # Of these, zipapp + entrypoint-module is the only combination where we can't both satisfy being
+    # able to re-exec a PEX based on argv[0] and support pickling.
+
+    unzip_env = make_env(PEX_UNZIP=1)
+
+    with_ep_pex = os.path.join(str(tmpdir), "test_with_ep.pex")
+    run_pex_command(
+        args=["-D", src_dir, "-e", "issues_1018:main", "-o", with_ep_pex]
+    ).assert_success()
+    assert expected_output == subprocess.check_output(args=[with_ep_pex], env=unzip_env)
+    assert expected_output == subprocess.check_output(args=[with_ep_pex])
+
+    no_ep_pex = os.path.join(str(tmpdir), "test_no_ep.pex")
+    run_pex_command(args=["-D", src_dir, "-o", no_ep_pex]).assert_success()
+    assert expected_output == subprocess.check_output(
+        args=[no_ep_pex, "-m", "issues_1018"], env=unzip_env
+    )
+    assert expected_output == subprocess.check_output(args=[no_ep_pex, "-m", "issues_1018"])
+
+    with_module_pex = os.path.join(str(tmpdir), "test_with_module.pex")
+    run_pex_command(
+        args=["-D", src_dir, "-m", "issues_1018", "-o", with_module_pex]
+    ).assert_success()
+    assert expected_output == subprocess.check_output(args=[with_module_pex], env=unzip_env)
+
+    # For the case of a PEX zip with a module entrypoint we cannot both get pickling working and
+    # support re-execution of the PEX file using sys.argv[0]. Prospective picklers need to either
+    # use a function entrypoint as in with_ep_pex or else re-structure their pickling to happen
+    # anywhere but in a __name__ == '__main__' module.
+    process = subprocess.Popen(args=[with_module_pex], stderr=subprocess.PIPE)
+    _, stderr = process.communicate()
+    assert process.returncode != 0
+    traceback_root = stderr.decode("utf-8").splitlines()[-1]
+    if IS_PYPY2:
+        assert "TypeError: can't pickle zipimporter objects" == traceback_root, traceback_root
+    else:
+        assert re.search(r"\bPicklingError\b", traceback_root) is not None, traceback_root
+        assert re.search(r"\b__main__\b", traceback_root) is not None
+        assert re.search(r"\badd\b", traceback_root) is not None
+        assert (
+            re.search(r"<function add at 0x[a-f0-9]+>", traceback_root) is not None
+        ), traceback_root
+
+    with_ep_venv_pex = os.path.join(str(tmpdir), "test_with_ep_venv.pex")
+    run_pex_command(
+        args=["-D", src_dir, "-e", "issues_1018:main", "-o", with_ep_venv_pex, "--venv"]
+    ).assert_success()
+    assert expected_output == subprocess.check_output(args=[with_ep_venv_pex])
+
+    no_ep_venv_pex = os.path.join(str(tmpdir), "test_no_ep_venv.pex")
+    result = run_pex_command(args=["-D", src_dir, "-o", no_ep_venv_pex, "--venv", "--seed"])
+    result.assert_success()
+    no_ep_venv_pex_bin = result.output.strip()
+    assert expected_output == subprocess.check_output(
+        args=[no_ep_venv_pex_bin, "-m", "issues_1018"]
+    )
+
+    with_module_venv_pex = os.path.join(str(tmpdir), "test_with_module_venv.pex")
+    run_pex_command(
+        args=["-D", src_dir, "-m", "issues_1018", "-o", with_module_venv_pex, "--venv"]
+    ).assert_success()
+    assert expected_output == subprocess.check_output(args=[with_module_venv_pex])
diff --git a/tests/test_pex.py b/tests/test_pex.py
--- a/tests/test_pex.py
+++ b/tests/test_pex.py
@@ -611,8 +611,17 @@ def test_execute_interpreter_dashm_module():
     with temporary_dir() as pex_chroot:
         pex_builder = PEXBuilder(path=pex_chroot)
         pex_builder.add_source(None, "foo/__init__.py")
-        with tempfile.NamedTemporaryFile() as fp:
-            fp.write(b'import sys; print(" ".join(sys.argv))')
+        with tempfile.NamedTemporaryFile(mode="w") as fp:
+            fp.write(
+                dedent(
+                    """\
+                    import os
+                    import sys
+
+                    print("{} {}".format(os.path.realpath(sys.argv[0]), " ".join(sys.argv[1:])))
+                    """
+                )
+            )
             fp.flush()
             pex_builder.add_source(fp.name, "foo/bar.py")
         pex_builder.freeze()
@@ -626,7 +635,9 @@ def test_execute_interpreter_dashm_module():
         stdout, stderr = process.communicate()
 
         assert 0 == process.returncode
-        assert b"foo.bar one two\n" == stdout
+        assert "{} one two\n".format(
+            os.path.realpath(os.path.join(pex_chroot, "foo/bar.py"))
+        ) == stdout.decode("utf-8")
         assert b"" == stderr