Repository: pex-tool/pex (Python)

Issue:
Dependencies specified in `extra_requires` don't work from 1.5.0 onwards
### Example:
`setup.py`:
```
from setuptools import setup, find_packages

setup(
    name = 'bad-pex-example',
    version = '0.1',
    python_requires = '>=3.5',
    install_requires = [],
    extras_require = { 'master': [ 'schema>=0.6' ] }
)
```

`bad/__init__.py`:
```
import schema

def entry():
    print('it works')
```

In version `1.4.9`, `pex -o bad.pex -e 'bad:entry' '.[master]'` produces a working executable. In all later versions (up to current latest `1.6.7`) the resulting executable fails to import the dependency:
```
Traceback (most recent call last):
  File ".bootstrap/pex/pex.py", line 352, in execute
  File ".bootstrap/pex/pex.py", line 284, in _wrap_coverage
  File ".bootstrap/pex/pex.py", line 315, in _wrap_profiling
  File ".bootstrap/pex/pex.py", line 397, in _execute
  File ".bootstrap/pex/pex.py", line 495, in execute_entry
  File ".bootstrap/pex/pex.py", line 509, in execute_pkg_resources
  File ".bootstrap/pex/vendor/_vendored/setuptools/pkg_resources/__init__.py", line 2384, in resolve
  File "/home/david/work/bad-pex-example/bad/__init__.py", line 1, in <module>
    import schema
ModuleNotFoundError: No module named 'schema'
```
`pex` does still get the dependency and does something with it:
```
❯ pex -o bad.pex -e 'bad:entry' '.[master]' -vvvv
pex: Constructed RequestsContext context <pex.http.RequestsContext object at 0x7f83f8be4908>       
pex: Constructed RequestsContext context <pex.http.RequestsContext object at 0x7f83f8be4240>
pex: Constructed RequestsContext context <pex.http.RequestsContext object at 0x7f83f8b982b0>                                                                                         
pex: crawling link i=0 link=Link('file:///home/david/.pex/build') follow_links=False
pex: crawling link i=0 link=Link('https://pypi.org/simple/schema/') follow_links=False
pex: Building pex :: Resolving distributions :: Fetching file:///home/david/.pex/build/schema-0.7.0-py2.py3-none-any.whl :: Fetching file:///home/david/.pex/build/schema-0.7.0-py2.ppex: Building pex :: Resolving distributions :: Translating /tmp/tmpa2mbu1dq/schema-0.7.0-py2.py3-none-any.whl into distribution                                                     pex: Constructed RequestsContext context <pex.http.RequestsContext object at 0x7f83f8374208>                                    
pex: crawling link i=0 link=Link('https://pypi.org/simple/contextlib2/') follow_links=False
pex: Building pex :: Resolving distributions :: Fetching file:///home/david/.pex/build/contextlib2-0.5.5-py2.py3-none-any.whl :: Fetching file:///home/david/.pex/build/contextlib2-0pex: Building pex :: Resolving distributions :: Translating /tmp/tmpf5_mo387/contextlib2-0.5.5-py2.py3-none-any.whl into distribution                                                  bad-pex-example==0.1 -> bad-pex-example 0.1
  schema>=0.6; extra == "master" -> schema 0.7.0
  contextlib2==0.5.5 -> contextlib2 0.5.5
pex: Building pex: 470.8ms                                                                                                           
pex:   Resolving interpreters: 0.0ms
pex:   Installing /tmp/tmp1zk6gp55: 209.2ms
pex:     Isolating pex in Chroot(/tmp/tmp6s9psqsi {fs:}): 15.1ms
pex:   Resolving distributions: 260.6ms
pex:     Fetching file:///tmp/tmp1zk6gp55/bad-pex-example-0.1.tar.gz: 0.5ms
pex:       Fetching file:///tmp/tmp1zk6gp55/bad-pex-example-0.1.tar.gz: 0.1ms
pex:     Translating /tmp/tmppx_oldwr/bad-pex-example-0.1.tar.gz into distribution: 169.7ms
pex:       Packaging bad-pex-example: 168.0ms
pex:         Installing /tmp/tmpuosc6ijy: 166.6ms
pex:     Fetching file:///home/david/.pex/build/schema-0.7.0-py2.py3-none-any.whl: 0.4ms
pex:       Fetching file:///home/david/.pex/build/schema-0.7.0-py2.py3-none-any.whl: 0.2ms
pex:     Translating /tmp/tmpa2mbu1dq/schema-0.7.0-py2.py3-none-any.whl into distribution: 0.8ms
pex:     Fetching file:///home/david/.pex/build/contextlib2-0.5.5-py2.py3-none-any.whl: 1.2ms
pex:       Fetching file:///home/david/.pex/build/contextlib2-0.5.5-py2.py3-none-any.whl: 0.4ms
pex:     Translating /tmp/tmpf5_mo387/contextlib2-0.5.5-py2.py3-none-any.whl into distribution: 1.9ms
pex: Vendoring pkg_resources from VendorSpec(key='setuptools', version='40.6.2') @ /home/david/.local/lib/python3.6/site-packages/pex/vendor/_vendored/setuptools
Saving PEX file to bad.pex
```
and `schema` still does end up in `~/.pex/install` after running the executable:
```
❯ ls -lah ~/.pex/install
total 20K
drwxrwxr-x 5 david david 4.0K Jun 17 11:42 .
drwxrwxr-x 3 david david 4.0K Jun 17 11:42 ..
drwxrwxr-x 4 david david 4.0K Jun 17 11:42 bad_pex_example-0.1-py3-none-any.whl.09110f078728861fa50039d90343b57aebf726ca
drwxrwxr-x 4 david david 4.0K Jun 17 11:42 contextlib2-0.5.5-py2.py3-none-any.whl.edad6efe1aebf29bea2d8a9e8dab93fbbc2555bc
drwxrwxr-x 4 david david 4.0K Jun 17 11:42 schema-0.7.0-py2.py3-none-any.whl.5b7491a4df1576f5f2912f4158d096cd06d90199
```
Many thanks!


Discussion:
I repro after:
```diff
$ diff -u /tmp/setup.py setup.py 
--- /tmp/setup.py	2019-06-22 15:38:38.599110687 -0600
+++ setup.py	2019-06-22 15:08:11.585058403 -0600
@@ -4,6 +4,7 @@
     name = 'bad-pex-example',
     version = '0.1',
     python_requires = '>=3.5',
+    packages = ['bad'],
     install_requires = [],
     extras_require = { 'master': [ 'schema>=0.6' ] }
 )
```

The bug was introduced in #582 here:
https://github.com/pantsbuild/pex/blob/fb35498778bf722a33133ed351cea5b47a523c39/pex/resolver.py#L327

The requirement string returned by a distribution for itself will never contain extras that were asked for when installing it. We have that information though in the resolvable the distribution was created from, so this is a one line fix with more lines to add a test.



Patch:
diff --git a/pex/resolver.py b/pex/resolver.py
index 3fd9a764d..f5b285121 100644
--- a/pex/resolver.py
+++ b/pex/resolver.py
@@ -20,7 +20,7 @@
 from pex.platforms import Platform
 from pex.resolvable import ResolvableRequirement, resolvables_from_iterable
 from pex.resolver_options import ResolverOptionsBuilder
-from pex.third_party.pkg_resources import safe_name
+from pex.third_party.pkg_resources import Distribution, Requirement, safe_name
 from pex.tracer import TRACER
 from pex.util import DistributionHelper
 
@@ -158,9 +158,14 @@ def map_packages(resolved_packages):
     return _ResolvableSet([map_packages(rp) for rp in self.__tuples])
 
 
-class ResolvedDistribution(namedtuple('ResolvedDistribution', 'requirement distribution')):
+class ResolvedDistribution(namedtuple('ResolvedDistribution', ['requirement', 'distribution'])):
   """A requirement and the resolved distribution that satisfies it."""
 
+  def __new__(cls, requirement, distribution):
+    assert isinstance(requirement, Requirement)
+    assert isinstance(distribution, Distribution)
+    return super(ResolvedDistribution, cls).__new__(cls, requirement, distribution)
+
 
 class Resolver(object):
   """Interface for resolving resolvable entities into python packages."""
@@ -333,6 +338,7 @@ def resolve(self, resolvables, resolvable_set=None):
         requirement = resolvable.requirement
       else:
         requirement = distribution.as_requirement()
+        requirement.extras = tuple(resolvable.extras())
       dists.append(ResolvedDistribution(requirement=requirement,
                                         distribution=distribution))
     return dists


Test patch:
diff --git a/pex/testing.py b/pex/testing.py
index c3af78b1d..bc2486bd1 100644
--- a/pex/testing.py
+++ b/pex/testing.py
@@ -99,30 +99,47 @@ def write_zipfile(directory, dest, reverse=False):
   return dest
 
 
-PROJECT_CONTENT = {
-  'setup.py': dedent('''
+@contextlib.contextmanager
+def make_project(name='my_project',
+                 version='0.0.0',
+                 zip_safe=True,
+                 install_reqs=None,
+                 extras_require=None):
+
+  project_content = {
+    'setup.py': dedent('''
       from setuptools import setup
 
       setup(
-          name=%(project_name)r,
-          version=%(version)r,
-          zip_safe=%(zip_safe)r,
-          packages=['my_package'],
-          scripts=[
-              'scripts/hello_world',
-              'scripts/shell_script',
-          ],
-          package_data={'my_package': ['package_data/*.dat']},
-          install_requires=%(install_requires)r,
+        name=%(project_name)r,
+        version=%(version)r,
+        zip_safe=%(zip_safe)r,
+        packages=[%(project_name)r],
+        scripts=[
+          'scripts/hello_world',
+          'scripts/shell_script',
+        ],
+        package_data={%(project_name)r: ['package_data/*.dat']},
+        install_requires=%(install_requires)r,
+        extras_require=%(extras_require)r,
       )
-  '''),
-  'scripts/hello_world': '#!/usr/bin/env python\nprint("hello world!")\n',
-  'scripts/shell_script': '#!/usr/bin/env bash\necho hello world\n',
-  'my_package/__init__.py': 0,
-  'my_package/my_module.py': 'def do_something():\n  print("hello world!")\n',
-  'my_package/package_data/resource1.dat': 1000,
-  'my_package/package_data/resource2.dat': 1000,
-}
+    '''),
+    'scripts/hello_world': '#!/usr/bin/env python\nprint("hello world!")\n',
+    'scripts/shell_script': '#!/usr/bin/env bash\necho hello world\n',
+    os.path.join(name, '__init__.py'): 0,
+    os.path.join(name, 'my_module.py'): 'def do_something():\n  print("hello world!")\n',
+    os.path.join(name, 'package_data/resource1.dat'): 1000,
+    os.path.join(name, 'package_data/resource2.dat'): 1000,
+  }
+
+  interp = {'project_name': name,
+            'version': version,
+            'zip_safe': zip_safe,
+            'install_requires': install_reqs or [],
+            'extras_require': extras_require or {}}
+
+  with temporary_content(project_content, interp=interp) as td:
+    yield td
 
 
 @contextlib.contextmanager
@@ -131,29 +148,38 @@ def make_installer(name='my_project',
                    installer_impl=EggInstaller,
                    zip_safe=True,
                    install_reqs=None,
+                   extras_require=None,
                    interpreter=None,
                    **kwargs):
-  interp = {'project_name': name,
-            'version': version,
-            'zip_safe': zip_safe,
-            'install_requires': install_reqs or []}
-  with temporary_content(PROJECT_CONTENT, interp=interp) as td:
+
+  with make_project(name=name,
+                    version=version,
+                    zip_safe=zip_safe,
+                    install_reqs=install_reqs,
+                    extras_require=extras_require) as td:
     yield installer_impl(td, interpreter=interpreter, **kwargs)
 
 
 @contextlib.contextmanager
-def make_source_dir(name='my_project', version='0.0.0', install_reqs=None):
-  interp = {'project_name': name,
-            'version': version,
-            'zip_safe': True,
-            'install_requires': install_reqs or []}
-  with temporary_content(PROJECT_CONTENT, interp=interp) as td:
+def make_source_dir(name='my_project', version='0.0.0', install_reqs=None, extras_require=None):
+  with make_project(name=name,
+                    version=version,
+                    install_reqs=install_reqs,
+                    extras_require=extras_require) as td:
     yield td
 
 
-def make_sdist(name='my_project', version='0.0.0', zip_safe=True, install_reqs=None):
-  with make_installer(name=name, version=version, installer_impl=Packager, zip_safe=zip_safe,
-                      install_reqs=install_reqs) as packager:
+def make_sdist(name='my_project',
+               version='0.0.0',
+               zip_safe=True,
+               install_reqs=None,
+               extras_require=None):
+  with make_installer(name=name,
+                      version=version,
+                      installer_impl=Packager,
+                      zip_safe=zip_safe,
+                      install_reqs=install_reqs,
+                      extras_require=extras_require) as packager:
     return packager.sdist()
 
 
diff --git a/tests/test_integration.py b/tests/test_integration.py
index f3be9331f..d9a336419 100644
--- a/tests/test_integration.py
+++ b/tests/test_integration.py
@@ -13,7 +13,7 @@
 
 import pytest
 
-from pex.common import safe_sleep
+from pex.common import safe_copy, safe_sleep
 from pex.compatibility import WINDOWS, nested, to_bytes
 from pex.installer import EggInstaller
 from pex.pex_info import PexInfo
@@ -30,6 +30,8 @@
     PY36,
     ensure_python_interpreter,
     get_dep_dist_names_from_pex,
+    make_sdist,
+    make_source_dir,
     run_pex_command,
     run_simple_pex,
     run_simple_pex_test,
@@ -1461,3 +1463,22 @@ def test_reproducible_build_python_flag():
 
 def test_reproducible_build_python_shebang_flag():
   assert_reproducible_build(['--python-shebang=/usr/bin/python'])
+
+
+def test_issues_736_requirement_setup_py_with_extras():
+  with make_source_dir(name='project1',
+                       version='1.0.0',
+                       extras_require={'foo': ['project2']}) as project1_dir:
+    project2_sdist = make_sdist(name='project2', version='2.0.0')
+    with temporary_dir() as td:
+      safe_copy(project2_sdist, os.path.join(td, os.path.basename(project2_sdist)))
+
+      project1_pex = os.path.join(td, 'project1.pex')
+      result = run_pex_command(['-f', td, '-o', project1_pex, '{}[foo]'.format(project1_dir)])
+      result.assert_success()
+
+      output = subprocess.check_output(
+        [project1_pex, '-c', 'from project2 import my_module; my_module.do_something()'],
+        env=make_env(PEX_INTERPRETER='1')
+      )
+      assert output.decode('utf-8').strip() == u'hello world!'
diff --git a/tests/test_pex_builder.py b/tests/test_pex_builder.py
index 19a583bc9..8b0624614 100644
--- a/tests/test_pex_builder.py
+++ b/tests/test_pex_builder.py
@@ -17,7 +17,7 @@
 
 exe_main = """
 import sys
-from my_package.my_module import do_something
+from p1.my_module import do_something
 do_something()
 
 with open(sys.argv[1], 'w') as fp:
@@ -27,7 +27,7 @@
 wheeldeps_exe_main = """
 import sys
 from pyparsing import *
-from my_package.my_module import do_something
+from p1.my_module import do_something
 do_something()
 
 with open(sys.argv[1], 'w') as fp:
diff --git a/tests/test_resolver.py b/tests/test_resolver.py
index 2ca79dcea..c7f60e8b6 100644
--- a/tests/test_resolver.py
+++ b/tests/test_resolver.py
@@ -13,7 +13,8 @@
 from pex.resolvable import ResolvableRequirement
 from pex.resolver import Resolver, Unsatisfiable, _ResolvableSet, resolve_multi
 from pex.resolver_options import ResolverOptionsBuilder
-from pex.testing import make_sdist, temporary_dir
+from pex.testing import make_sdist, make_source_dir, temporary_dir
+from pex.third_party.pkg_resources import Requirement
 
 
 def do_resolve_multi(*args, **kwargs):
@@ -337,3 +338,37 @@ def test_resolvable_set_built():
   updated_rs.merge(rq, [binary_pkg])
   assert updated_rs.get('foo') == set([binary_pkg])
   assert updated_rs.packages() == [(rq, set([binary_pkg]), None, False)]
+
+
+def _parse_requirement(req):
+  return Requirement.parse(str(req))
+
+
+def test_resolve_extra_setup_py():
+  with make_source_dir(name='project1',
+                       version='1.0.0',
+                       extras_require={'foo': ['project2']}) as project1_dir:
+    project2_sdist = make_sdist(name='project2', version='2.0.0')
+    with temporary_dir() as td:
+      safe_copy(project2_sdist, os.path.join(td, os.path.basename(project2_sdist)))
+      fetchers = [Fetcher([td])]
+
+      resolved_dists = do_resolve_multi(['{}[foo]'.format(project1_dir)], fetchers=fetchers)
+      assert ({_parse_requirement(req) for req in ('project1[foo]==1.0.0',
+                                                   'project2; extra=="foo"')} ==
+              {_parse_requirement(resolved_dist.requirement) for resolved_dist in resolved_dists})
+
+
+def test_resolve_extra_sdist():
+  project1_sdist = make_sdist(name='project1',
+                              version='1.0.0',
+                              extras_require={'foo': ['project2']})
+  project2_sdist = make_sdist(name='project2', version='2.0.0')
+  with temporary_dir() as td:
+    for sdist in (project1_sdist, project2_sdist):
+      safe_copy(sdist, os.path.join(td, os.path.basename(sdist)))
+    fetchers = [Fetcher([td])]
+
+    resolved_dists = do_resolve_multi(['project1[foo]'], fetchers=fetchers)
+    assert ({_parse_requirement(req) for req in ('project1[foo]', 'project2; extra=="foo"')} ==
+            {_parse_requirement(resolved_dist.requirement) for resolved_dist in resolved_dists})