Extra requirement support (extras_require) (#1363)

piptools/scripts/compile.py

@@ -2,12 +2,13 @@
 import shlex
 import sys
 import tempfile
-from typing import Any, BinaryIO, Optional, Tuple, cast
+from typing import Any, BinaryIO, List, Optional, Tuple, cast
 
 import click
 from click.utils import LazyFile, safecall
 from pep517 import meta
 from pip._internal.commands import create_command
+from pip._internal.req import InstallRequirement
 from pip._internal.req.constructors import install_req_from_line
 from pip._internal.utils.misc import redact_auth_from_url
 
@@ -19,7 +20,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,
+    drop_extras,
+    is_pinned_requirement,
+    key_from_ireq,
+)
 from ..writer import OutputWriter
 
 DEFAULT_REQUIREMENTS_FILE = "requirements.in"
@@ -61,6 +68,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 +219,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,
@@ -336,7 +350,8 @@ def cli(
     # Parsing/collecting initial requirements
     ###
 
-    constraints = []
+    constraints: List[InstallRequirement] = []
+    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 +375,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 +394,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 +407,9 @@ 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.match_markers(extras)]
+    for req in constraints:
+        drop_extras(req)
 
     log.debug("Using indexes:")
     with log.indentation():

piptools/utils.py

@@ -6,6 +6,7 @@
     Dict,
     Iterable,
     Iterator,
+    List,
     Optional,
     Set,
     Tuple,
@@ -210,6 +211,57 @@ def dedup(iterable: Iterable[_T]) -> Iterable[_T]:
     return iter(dict.fromkeys(iterable))
 
 
+def drop_extras(ireq: InstallRequirement) -> None:
+    """Remove "extra" markers (PEP-508) from requirement."""
+    if ireq.markers is None:
+        return
+    ireq.markers._markers = _drop_extras(ireq.markers._markers)
+    if not ireq.markers._markers:
+        ireq.markers = None
+
+
+def _drop_extras(markers: List[_T]) -> List[_T]:
+    # drop `extra` tokens
+    to_remove: List[int] = []
+    for i, token in enumerate(markers):
+        # operator (and/or)
+        if isinstance(token, str):
+            continue
+        # sub-expression (inside braces)
+        if isinstance(token, list):
+            markers[i] = _drop_extras(token)  # type: ignore
+            if markers[i]:
+                continue
+            to_remove.append(i)
+            continue
+        # test expression (like `extra == "dev"`)
+        assert isinstance(token, tuple)
+        if token[0].value == "extra":
+            to_remove.append(i)
+    for i in reversed(to_remove):
+        markers.pop(i)
+
+    # drop duplicate bool operators (and/or)
+    to_remove = []
+    for i, (token1, token2) in enumerate(zip(markers, markers[1:])):
+        if not isinstance(token1, str):
+            continue
+        if not isinstance(token2, str):
+            continue
+        if token1 == "and":
+            to_remove.append(i)
+        else:
+            to_remove.append(i + 1)
+    for i in reversed(to_remove):
+        markers.pop(i)
+    if markers and isinstance(markers[0], str):
+        markers.pop(0)
+    if markers and isinstance(markers[-1], str):
+        markers.pop(-1)
+
+    return markers
+
+
 def get_hashes_from_ireq(ireq: InstallRequirement) -> Set[str]:
     """
     Given an InstallRequirement, return a set of string hashes in the format

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