Support scripts with inline script metadata as input files

Author: chrysle

README.md

@@ -32,7 +32,8 @@ project's virtual environment.
 
 The `pip-compile` command lets you compile a `requirements.txt` file from
 your dependencies, specified in either `pyproject.toml`, `setup.cfg`,
-`setup.py`, or `requirements.in`.
+`setup.py`, `requirements.in`, or pure-Python scripts containing
+[inline script metadata](https://packaging.python.org/en/latest/specifications/inline-script-metadata/).
 
 Run it with `pip-compile` or `python -m piptools compile` (or
 `pipx run --spec pip-tools pip-compile` if `pipx` was installed with the

piptools/scripts/compile.py

@@ -2,9 +2,11 @@
 
 import itertools
 import os
+import re
 import shlex
 import sys
 import tempfile
+import tomllib
 from pathlib import Path
 from typing import IO, Any, BinaryIO, cast
 
@@ -43,6 +45,10 @@
 DEFAULT_REQUIREMENTS_OUTPUT_FILE = "requirements.txt"
 METADATA_FILENAMES = frozenset({"setup.py", "setup.cfg", "pyproject.toml"})
 
+INLINE_SCRIPT_METADATA_REGEX = (
+    r"(?m)^# /// (?P<type>[a-zA-Z0-9-]+)$\s(?P<content>(^#(| .*)$\s)+)^# ///$"
+)
+
 
 def _determine_linesep(
     strategy: str = "preserve", filenames: tuple[str, ...] = ()
@@ -170,7 +176,8 @@ def cli(
 ) -> None:
     """
     Compiles requirements.txt from requirements.in, pyproject.toml, setup.cfg,
-    or setup.py specs.
+    or setup.py specs, as well as Python scripts containing inline script
+    metadata.
     """
     if color is not None:
         ctx.color = color
@@ -344,14 +351,50 @@ def cli(
             )
             raise click.BadParameter(msg)
 
-        if src_file == "-":
-            # pip requires filenames and not files. Since we want to support
-            # piping from stdin, we need to briefly save the input from stdin
-            # to a temporary file and have pip read that.  also used for
+        if src_file == "-" or (
+            os.path.basename(src_file).endswith(".py") and not is_setup_file
+        ):
+            # pip requires filenames and not files.  Since we want to support
+            # piping from stdin, and inline script metadadat within Python
+            # scripts, we need to briefly save the input or extracted script
+            # dependencies to a temporary file and have pip read that.  Also used for
             # reading requirements from install_requires in setup.py.
+            if os.path.basename(src_file).endswith(".py"):
+                # Probably contains inline script metadata
+                with open(src_file, encoding="utf-8") as f:
+                    script = f.read()
+                name = "script"
+                matches = list(
+                    filter(
+                        lambda m: m.group("type") == name,
+                        re.finditer(INLINE_SCRIPT_METADATA_REGEX, script),
+                    )
+                )
+                if len(matches) > 1:
+                    raise ValueError(f"Multiple {name} blocks found")
+                elif len(matches) == 1:
+                    content = "".join(
+                        line[2:] if line.startswith("# ") else line[1:]
+                        for line in matches[0]
+                        .group("content")
+                        .splitlines(keepends=True)
+                    )
+                    metadata = tomllib.loads(content)
+                    reqs_str = metadata.get("dependencies", [])
+                    tmpfile = tempfile.NamedTemporaryFile(mode="wt", delete=False)
+                    input_reqs = "\n".join(reqs_str)
+                    comes_from = (
+                        f"{os.path.basename(src_file)} (inline script metadata)"
+                    )
+                else:
+                    raise PipToolsError(
+                        "Input script does not contain valid inline script metadata!"
+                    )
+            else:
+                input_reqs = sys.stdin.read()
+                comes_from = "-r -"
             tmpfile = tempfile.NamedTemporaryFile(mode="wt", delete=False)
-            tmpfile.write(sys.stdin.read())
-            comes_from = "-r -"
+            tmpfile.write(input_reqs)
             tmpfile.flush()
             reqs = list(
                 parse_requirements(

tests/test_cli_compile.py

@@ -17,6 +17,7 @@
 from pip._vendor.packaging.version import Version
 
 from piptools.build import ProjectMetadata
+from piptools.exceptions import PipToolsError
 from piptools.scripts.compile import cli
 from piptools.utils import (
     COMPILE_EXCLUDE_OPTIONS,
@@ -3771,3 +3772,70 @@ def test_stdout_should_not_be_read_when_stdin_is_not_a_plain_file(
     out = runner.invoke(cli, [req_in.as_posix(), "--output-file", fifo.as_posix()])
 
     assert out.exit_code == 0, out
+
+
+def test_compile_inline_script_metadata(runner, tmp_path, current_resolver):
+    (tmp_path / "script.py").write_text(
+        dedent(
+            """
+            # /// script
+            # dependencies = [
+            #   "small-fake-with-deps",
+            # ]
+            # ///
+            """
+        )
+    )
+    out = runner.invoke(
+        cli,
+        [
+            "--no-build-isolation",
+            "--no-header",
+            "--no-emit-options",
+            "--find-links",
+            os.fspath(MINIMAL_WHEELS_PATH),
+            os.fspath(tmp_path / "script.py"),
+            "--output-file",
+            "-",
+        ],
+    )
+    expected = r"""small-fake-a==0.1
+    # via small-fake-with-deps
+small-fake-with-deps==0.1
+    # via script.py (inline script metadata)
+"""
+    assert out.exit_code == 0
+    assert expected == out.stdout
+
+
+def test_compile_inline_script_metadata_invalid(runner, tmp_path, current_resolver):
+    (tmp_path / "script.py").write_text(
+        dedent(
+            """
+            # /// invalid-name
+            # dependencies = [
+            #   "small-fake-a",
+            #   "small-fake-b",
+            # ]
+            # ///
+            """
+        )
+    )
+    with pytest.raises(
+        PipToolsError, match="does not contain valid inline script metadata"
+    ):
+        runner.invoke(
+            cli,
+            [
+                "--no-build-isolation",
+                "--no-header",
+                "--no-annotate",
+                "--no-emit-options",
+                "--find-links",
+                os.fspath(MINIMAL_WHEELS_PATH),
+                os.fspath(tmp_path / "script.py"),
+                "--output-file",
+                "-",
+            ],
+            catch_exceptions=False,
+        )