Extra requirement support (extras_require)

piptools/scripts/compile.py

@@ -19,7 +19,13 @@
 from ..repositories import LocalRequirementsRepository, PyPIRepository
 from ..repositories.base import BaseRepository
 from ..resolver import Resolver
-from ..utils import UNSAFE_PACKAGES, dedup, is_pinned_requirement, key_from_ireq
+from ..utils import (
+    UNSAFE_PACKAGES,
+    dedup,
+    is_pinned_requirement,
+    key_from_ireq,
+    req_check_markers,
+)
 from ..writer import OutputWriter
 
 DEFAULT_REQUIREMENTS_FILE = "requirements.in"
@@ -61,6 +67,12 @@ def _get_default_option(option_name: str) -> Any:
     is_flag=True,
     help="Clear any caches upfront, rebuild from scratch",
 )
+@click.option(
+    "--extra",
+    "extras",
+    multiple=True,
+    help="Names of extras_require to install",
+)
 @click.option(
     "-f",
     "--find-links",
@@ -206,22 +218,23 @@ def cli(
     dry_run: bool,
     pre: bool,
     rebuild: bool,
-    find_links: Tuple[str],
+    extras: Tuple[str, ...],
+    find_links: Tuple[str, ...],
     index_url: str,
-    extra_index_url: Tuple[str],
+    extra_index_url: Tuple[str, ...],
     cert: Optional[str],
     client_cert: Optional[str],
-    trusted_host: Tuple[str],
+    trusted_host: Tuple[str, ...],
     header: bool,
     emit_trusted_host: bool,
     annotate: bool,
     upgrade: bool,
-    upgrade_packages: Tuple[str],
+    upgrade_packages: Tuple[str, ...],
     output_file: Optional[LazyFile],
     allow_unsafe: bool,
     generate_hashes: bool,
     reuse_hashes: bool,
-    src_files: Tuple[str],
+    src_files: Tuple[str, ...],
     max_rounds: int,
     build_isolation: bool,
     emit_find_links: bool,
@@ -337,6 +350,7 @@ def cli(
     ###
 
     constraints = []
+    setup_file_found = False
     for src_file in src_files:
         is_setup_file = os.path.basename(src_file) in METADATA_FILENAMES
         if src_file == "-":
@@ -360,6 +374,7 @@ def cli(
                 req.comes_from = comes_from
             constraints.extend(reqs)
         elif is_setup_file:
+            setup_file_found = True
             dist = meta.load(os.path.dirname(os.path.abspath(src_file)))
             comes_from = f"{dist.metadata.get_all('Name')[0]} ({src_file})"
             constraints.extend(
@@ -378,6 +393,10 @@ def cli(
                 )
             )
 
+    if extras and not setup_file_found:
+        msg = "--extra has effect only with setup.py and PEP-517 input formats"
+        raise click.BadParameter(msg)
+
     primary_packages = {
         key_from_ireq(ireq) for ireq in constraints if not ireq.constraint
     }
@@ -387,16 +406,7 @@ def cli(
         ireq for key, ireq in upgrade_install_reqs.items() if key in allowed_upgrades
     )
 
-    # Filter out pip environment markers which do not match (PEP496)
-    constraints = [
-        req
-        for req in constraints
-        if req.markers is None
-        # We explicitly set extra=None to filter out optional requirements
-        # since evaluating an extra marker with no environment raises UndefinedEnvironmentName
-        # (see https://packaging.pypa.io/en/latest/markers.html#usage)
-        or req.markers.evaluate({"extra": None})
-    ]
+    constraints = [req for req in constraints if req_check_markers(req, extras=extras)]
 
     log.debug("Using indexes:")
     with log.indentation():

piptools/utils.py

@@ -163,6 +163,19 @@ def is_pinned_requirement(ireq: InstallRequirement) -> bool:
     return spec.operator in {"==", "==="} and not spec.version.endswith(".*")
 
 
+def req_check_markers(ireq: InstallRequirement, extras: Tuple[str, ...]) -> bool:
+    """
+    1. Check if the environment markers match (PEP-496).
+    2. Check if the requirement isn't extra or is included into extras to install.
+    """
+    if not ireq.markers or ireq.markers.evaluate({"extra": None}):
+        return True
+    for extra in extras:
+        if ireq.markers.evaluate({"extra": extra}):
+            return True
+    return False
+
+
 def as_tuple(ireq: InstallRequirement) -> Tuple[str, str, Tuple[str, ...]]:
     """
     Pulls out the (name: str, version:str, extras:(str)) tuple from

tests/conftest.py

@@ -322,3 +322,42 @@ def _make_sdist(package_dir, dist_dir, *args):
         return run_setup_file(package_dir, "sdist", "--dist-dir", str(dist_dir), *args)
 
     return _make_sdist
+
+
+@pytest.fixture
+def make_module(tmpdir):
+    """
+    Make a metadata file with the given name and content and a fake module.
+    """
+
+    def _make_module(fname, content):
+        path = os.path.join(tmpdir, "sample_lib")
+        os.mkdir(path)
+        path = os.path.join(tmpdir, "sample_lib", "__init__.py")
+        with open(path, "w") as stream:
+            stream.write("'example module'\n__version__ = '1.2.3'")
+        path = os.path.join(tmpdir, fname)
+        with open(path, "w") as stream:
+            stream.write(dedent(content))
+        return path
+
+    return _make_module
+
+
+@pytest.fixture
+def fake_dists(tmpdir, make_package, make_wheel):
+    """
+    Generate distribution packages `small-fake-{a..f}`
+    """
+    dists_path = os.path.join(tmpdir, "dists")
+    pkgs = [
+        make_package("small-fake-a", version="0.1"),
+        make_package("small-fake-b", version="0.2"),
+        make_package("small-fake-c", version="0.3"),
+        make_package("small-fake-d", version="0.4"),
+        make_package("small-fake-e", version="0.5"),
+        make_package("small-fake-f", version="0.6"),
+    ]
+    for pkg in pkgs:
+        make_wheel(pkg, dists_path)
+    return dists_path