Add --only-extra option to pip-compile

piptools/scripts/compile.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import itertools
 import os
 import shlex
 import sys
@@ -82,6 +81,7 @@ def _determine_linesep(
 @options.rebuild
 @options.extra
 @options.all_extras
+@options.only_extra
 @options.find_links
 @options.index_url
 @options.no_index
@@ -123,6 +123,7 @@ def cli(
     rebuild: bool,
     extras: tuple[str, ...],
     all_extras: bool,
+    only_extras: tuple[str, ...],
     find_links: tuple[str, ...],
     index_url: str,
     no_index: bool,
@@ -383,11 +384,26 @@ def cli(
             req.comes_from = None
         constraints.extend(reqs)
 
-    extras = tuple(itertools.chain.from_iterable(ex.split(",") for ex in extras))
+    extra_options_lookup = {
+        options.EXTRA_OPTION: extras,
+        options.ONLY_EXTRA_OPTION: only_extras,
+        options.ALL_EXTRAS_OPTION: all_extras,
+    }
+
+    if not setup_file_found:
+        for option, value in extra_options_lookup.items():
+            if value:
+                raise click.BadParameter(
+                    f"{option} has effect only with setup.py and PEP-517 input formats"
+                )
 
-    if extras and not setup_file_found:
-        msg = "--extra has effect only with setup.py and PEP-517 input formats"
-        raise click.BadParameter(msg)
+    if only_extras:
+        # Note, the order is important here because --all-extra flag populates extras
+        for option in (options.ALL_EXTRAS_OPTION, options.EXTRA_OPTION):
+            if extra_options_lookup[option]:
+                raise click.BadParameter(
+                    f"{option} and --only-extra are mutually exclusive."
+                )
 
     primary_packages = {
         key_from_ireq(ireq) for ireq in constraints if not ireq.constraint
@@ -397,7 +413,15 @@ def cli(
         ireq for key, ireq in upgrade_install_reqs.items() if key in primary_packages
     )
 
-    constraints = [req for req in constraints if req.match_markers(extras)]
+    if only_extras:
+        constraints = [
+            req
+            for req in constraints
+            if req.markers is not None and req.match_markers(only_extras)
+        ]
+    else:
+        constraints = [req for req in constraints if req.match_markers(extras)]
+
     for req in constraints:
         drop_extras(req)
 

piptools/scripts/options.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import itertools
 from typing import Any
 
 import click
@@ -9,6 +10,10 @@
 from piptools.locations import CACHE_DIR, DEFAULT_CONFIG_FILE_NAMES
 from piptools.utils import UNSAFE_PACKAGES, override_defaults_from_config_file
 
+EXTRA_OPTION = "--extra"
+ALL_EXTRAS_OPTION = "--all-extras"
+ONLY_EXTRA_OPTION = "--only-extra"
+
 
 def _get_default_option(option_name: str) -> Any:
     """
@@ -20,6 +25,12 @@ def _get_default_option(option_name: str) -> Any:
     return getattr(default_values, option_name)
 
 
+def _flatten_comma_separated_values(
+    ctx: click.Context, param: click.Parameter, value: tuple[str, ...]
+) -> tuple[str, ...]:
+    return tuple(itertools.chain.from_iterable(v.split(",") for v in value))
+
+
 help_option_names = ("-h", "--help")
 
 # The options used by pip-compile and pip-sync are presented in no specific order.
@@ -62,14 +73,15 @@ def _get_default_option(option_name: str) -> Any:
 )
 
 extra = click.option(
-    "--extra",
+    EXTRA_OPTION,
     "extras",
     multiple=True,
+    callback=_flatten_comma_separated_values,
     help="Name of an extras_require group to install; may be used more than once",
 )
 
 all_extras = click.option(
-    "--all-extras",
+    ALL_EXTRAS_OPTION,
     is_flag=True,
     default=False,
     help="Install all extras_require groups",
@@ -364,3 +376,11 @@ def _get_default_option(option_name: str) -> Any:
     is_flag=True,
     help="Restrict attention to user directory",
 )
+
+only_extra = click.option(
+    ONLY_EXTRA_OPTION,
+    "only_extras",
+    multiple=True,
+    callback=_flatten_comma_separated_values,
+    help="Name of the only extras_require group to install; may be used more than once",
+)

tests/test_cli_compile.py

@@ -2710,6 +2710,108 @@ def test_all_extras(fake_dists, runner, make_module, fname, content):
     )
 
 
+@pytest.mark.network
+@pytest.mark.parametrize(("fname", "content"), METADATA_TEST_CASES)
+def test_only_extras(fake_dists, runner, make_module, fname, content):
+    """
+    Test passing `--only-extra` includes only the specified extra.
+    """
+    meta_path = make_module(fname=fname, content=content)
+    out = runner.invoke(
+        cli,
+        [
+            "--only-extra",
+            "dev",
+            "--output-file",
+            "-",
+            "--quiet",
+            "--find-links",
+            fake_dists,
+            "--no-annotate",
+            "--no-emit-options",
+            "--no-header",
+            "--no-build-isolation",
+            meta_path,
+        ],
+    )
+    assert out.exit_code == 0, out
+    assert out.stdout == "small-fake-b==0.2\n"
+
+
+@pytest.mark.network
+@pytest.mark.parametrize(
+    "extra_opts",
+    (
+        pytest.param(("--only-extra", "dev", "--only-extra", "test"), id="singular"),
+        pytest.param(("--only-extra", "dev,test"), id="comma-separated"),
+    ),
+)
+@pytest.mark.parametrize(("fname", "content"), METADATA_TEST_CASES)
+def test_multiple_only_extras(
+    fake_dists, runner, make_module, fname, content, extra_opts
+):
+    """
+    Test passing multiple `--only-extra` params.
+    """
+    meta_path = make_module(fname=fname, content=content)
+    out = runner.invoke(
+        cli,
+        [
+            *extra_opts,
+            "--output-file",
+            "-",
+            "--quiet",
+            "--find-links",
+            fake_dists,
+            "--no-annotate",
+            "--no-emit-options",
+            "--no-header",
+            "--no-build-isolation",
+            meta_path,
+        ],
+    )
+    assert out.exit_code == 0, out.stderr
+    assert out.stdout == dedent(
+        """\
+        small-fake-b==0.2
+        small-fake-c==0.3
+        """
+    )
+
+
+# This should not depend on the metadata format so testing all cases is wasteful
+@pytest.mark.parametrize(("fname", "content"), METADATA_TEST_CASES[:1])
+@pytest.mark.parametrize(
+    "extra_opts",
+    (
+        pytest.param(("--extra", "test"), id="extra and only_extra"),
+        pytest.param(("--all-extras",), id="all_extras and only_extra"),
+    ),
+)
+def test_raise_error_when_mutual_exclusive_extra_options_are_passed(
+    fake_dists, runner, tmp_path, make_module, fname, content, extra_opts
+):
+    meta_path = make_module(fname=fname, content=content)
+    out = runner.invoke(
+        cli,
+        [
+            *extra_opts,
+            "--only-extra",
+            "dev",
+            "--find-links",
+            fake_dists,
+            "--no-build-isolation",
+            meta_path,
+            "--output-file",
+            "-",
+        ],
+    )
+
+    assert out.exit_code == 2
+    expected = f"{extra_opts[0]} and --only-extra are mutually exclusive."
+    assert expected in out.stderr
+
+
 # This should not depend on the metadata format so testing all cases is wasteful
 @pytest.mark.parametrize(("fname", "content"), METADATA_TEST_CASES[:1])
 def test_all_extras_fail_with_extra(fake_dists, runner, make_module, fname, content):
@@ -2738,16 +2840,24 @@ def test_all_extras_fail_with_extra(fake_dists, runner, make_module, fname, cont
     assert exp in out.stderr
 
 
-def test_extras_fail_with_requirements_in(runner, tmpdir):
+@pytest.mark.parametrize(
+    "options",
+    (
+        pytest.param(["--extra", "something"], id="extra"),
+        pytest.param(["--only-extra", "something"], id="only_extra"),
+        pytest.param(["--all-extras"], id="all_extras"),
+    ),
+)
+def test_extras_fail_with_requirements_in(runner, tmpdir, options):
     """
-    Test that passing `--extra` with `requirements.in` input file fails.
+    Test that passing any extra options with `requirements.in` input file fails.
     """
     path = os.path.join(tmpdir, "requirements.in")
     with open(path, "w") as stream:
         stream.write("\n")
-    out = runner.invoke(cli, ["-n", "--extra", "something", path])
+    out = runner.invoke(cli, options + ["-n", path])
     assert out.exit_code == 2
-    exp = "--extra has effect only with setup.py and PEP-517 input formats"
+    exp = f"{options[0]} has effect only with setup.py and PEP-517 input formats"
     assert exp in out.stderr