Fix resolution of setup.py project extras.

The change in pex-tool#582 introduced a bug where requested extras were not
preserved for local `setup.py` resolvables. Add failing tests for this
and fix.

Fixes pex-tool#736

pex/resolver.py

@@ -333,6 +333,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

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,29 @@ 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(**kwargs):
+  with make_installer(installer_impl=Packager, **kwargs) as packager:
     return packager.sdist()
 
 

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!'

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:

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})