Add [python].resolves_to_constraints_file

src/python/pants/backend/python/goals/lockfile.py

@@ -40,11 +40,19 @@
     WrappedGenerateLockfile,
 )
 from pants.core.util_rules.lockfile_metadata import calculate_invalidation_digest
-from pants.engine.fs import CreateDigest, Digest, DigestContents, FileContent
+from pants.engine.fs import (
+    CreateDigest,
+    Digest,
+    DigestContents,
+    FileContent,
+    GlobMatchErrorBehavior,
+    MergeDigests,
+    PathGlobs,
+)
 from pants.engine.process import ProcessCacheScope, ProcessResult
 from pants.engine.rules import Get, MultiGet, collect_rules, rule
 from pants.engine.target import AllTargets
-from pants.engine.unions import UnionRule
+from pants.engine.unions import UnionMembership, UnionRule
 from pants.util.docutil import bin_name
 from pants.util.logging import LogLevel
 from pants.util.ordered_set import FrozenOrderedSet
@@ -154,16 +162,46 @@ async def generate_lockfile(
     generate_lockfiles_subsystem: GenerateLockfilesSubsystem,
     python_repos: PythonRepos,
     python_setup: PythonSetup,
+    union_membership: UnionMembership,
 ) -> GenerateLockfileResult:
     if req.use_pex:
         pip_args_file = "__pip_args.txt"
         pip_args_file_content = "\n".join(
             [f"--no-binary {pkg}" for pkg in python_setup.no_binary]
             + [f"--only-binary {pkg}" for pkg in python_setup.only_binary]
         )
-        pip_args_file_digest = await Get(
+        pip_args_digest_get = Get(
             Digest, CreateDigest([FileContent(pip_args_file, pip_args_file_content.encode())])
         )
+
+        all_tool_resolve_names = tuple(
+            sentinel.resolve_name for sentinel in union_membership.get(GenerateToolLockfileSentinel)
+        )
+        constraints_file_path = python_setup.resolves_to_constraints_file(
+            all_tool_resolve_names
+        ).get(req.resolve_name)
+        constraints_digest_get = Get(
+            Digest,
+            PathGlobs(
+                [constraints_file_path] if constraints_file_path else [],
+                glob_match_error_behavior=GlobMatchErrorBehavior.error,
+                description_of_origin=softwrap(
+                    f"""
+                    the option `[python].resolves_to_constraints_file` for the resolve
+                    '{req.resolve_name}'
+                    """
+                ),
+            ),
+        )
+        constraints_args = (
+            (f"--constraints={constraints_file_path}",) if constraints_file_path else ()
+        )
+
+        pip_args_digest, constraints_digest = await MultiGet(
+            pip_args_digest_get, constraints_digest_get
+        )
+        input_digest = await Get(Digest, MergeDigests([pip_args_digest, constraints_digest]))
+
         header_delimiter = "//"
         result = await Get(
             ProcessResult,
@@ -194,12 +232,13 @@ async def generate_lockfile(
                     "--indent=2",
                     "-r",
                     pip_args_file,
+                    *constraints_args,
                     *python_repos.pex_args,
                     *python_setup.manylinux_pex_args,
                     *req.interpreter_constraints.generate_pex_arg_list(),
                     *req.requirements,
                 ),
-                additional_input_digest=pip_args_file_digest,
+                additional_input_digest=input_digest,
                 output_files=("lock.json",),
                 description=f"Generate lockfile for {req.resolve_name}",
                 # Instead of caching lockfile generation with LMDB, we instead use the invalidation

src/python/pants/backend/python/goals/lockfile_test.py

@@ -26,12 +26,27 @@
 from pants.util.strutil import strip_prefix
 
 
-def _generate(*, rule_runner: RuleRunner, use_pex: bool) -> str:
+@pytest.fixture
+def rule_runner() -> RuleRunner:
+    rule_runner = RuleRunner(
+        rules=[
+            *lockfile_rules(),
+            *pex.rules(),
+            QueryRule(GenerateLockfileResult, [GeneratePythonLockfile]),
+        ]
+    )
+    rule_runner.set_options([], env_inherit=PYTHON_BOOTSTRAP_ENV)
+    return rule_runner
+
+
+def _generate(
+    *, rule_runner: RuleRunner, use_pex: bool, ansicolors_version: str = "==1.1.8"
+) -> str:
     result = rule_runner.request(
         GenerateLockfileResult,
         [
             GeneratePythonLockfile(
-                requirements=FrozenOrderedSet(["ansicolors==1.1.8"]),
+                requirements=FrozenOrderedSet([f"ansicolors{ansicolors_version}"]),
                 interpreter_constraints=InterpreterConstraints(),
                 resolve_name="test",
                 lockfile_dest="test.lock",
@@ -41,19 +56,32 @@ def _generate(*, rule_runner: RuleRunner, use_pex: bool) -> str:
     )
     digest_contents = rule_runner.request(DigestContents, [result.digest])
     assert len(digest_contents) == 1
-    return digest_contents[0].content.decode()
+    content = digest_contents[0].content.decode()
+    if not use_pex:
+        return content
 
-
-def test_poetry_lockfile_generation() -> None:
-    rule_runner = RuleRunner(
-        rules=[
-            *lockfile_rules(),
-            *pex.rules(),
-            QueryRule(GenerateLockfileResult, [GeneratePythonLockfile]),
-        ]
+    pex_header = dedent(
+        f"""\
+        // This lockfile was autogenerated by Pants. To regenerate, run:
+        //
+        //    ./pants generate-lockfiles --resolve=test
+        //
+        // --- BEGIN PANTS LOCKFILE METADATA: DO NOT EDIT OR REMOVE ---
+        // {{
+        //   "version": 2,
+        //   "valid_for_interpreter_constraints": [],
+        //   "generated_with_requirements": [
+        //     "ansicolors{ansicolors_version}"
+        //   ]
+        // }}
+        // --- END PANTS LOCKFILE METADATA ---
+        """
     )
-    rule_runner.set_options([], env_inherit=PYTHON_BOOTSTRAP_ENV)
+    assert content.startswith(pex_header)
+    return strip_prefix(content, pex_header)
 
+
+def test_poetry_lockfile_generation(rule_runner: RuleRunner) -> None:
     poetry_lock = _generate(rule_runner=rule_runner, use_pex=False)
     assert poetry_lock.startswith("# This lockfile was autogenerated by Pants.")
     assert poetry_lock.rstrip().endswith(
@@ -69,41 +97,17 @@ def test_poetry_lockfile_generation() -> None:
 @pytest.mark.parametrize(
     ("no_binary", "only_binary"), ((False, False), (False, True), (True, False))
 )
-def test_pex_lockfile_generation(no_binary: bool, only_binary: bool) -> None:
-    rule_runner = RuleRunner(
-        rules=[
-            *lockfile_rules(),
-            *pex.rules(),
-            QueryRule(GenerateLockfileResult, [GeneratePythonLockfile]),
-        ]
-    )
+def test_pex_lockfile_generation(
+    rule_runner: RuleRunner, no_binary: bool, only_binary: bool
+) -> None:
     args = []
     if no_binary:
         args.append("--python-no-binary=ansicolors")
     if only_binary:
         args.append("--python-only-binary=ansicolors")
     rule_runner.set_options(args, env_inherit=PYTHON_BOOTSTRAP_ENV)
 
-    pex_header = dedent(
-        """\
-        // This lockfile was autogenerated by Pants. To regenerate, run:
-        //
-        //    ./pants generate-lockfiles --resolve=test
-        //
-        // --- BEGIN PANTS LOCKFILE METADATA: DO NOT EDIT OR REMOVE ---
-        // {
-        //   "version": 2,
-        //   "valid_for_interpreter_constraints": [],
-        //   "generated_with_requirements": [
-        //     "ansicolors==1.1.8"
-        //   ]
-        // }
-        // --- END PANTS LOCKFILE METADATA ---
-        """
-    )
-    pex_lock = _generate(rule_runner=rule_runner, use_pex=True)
-    assert pex_lock.startswith(pex_header)
-    lock_entry = json.loads(strip_prefix(pex_lock, pex_header))
+    lock_entry = json.loads(_generate(rule_runner=rule_runner, use_pex=True))
     reqs = lock_entry["locked_resolves"][0]["locked_requirements"]
     assert len(reqs) == 1
     assert reqs[0]["project_name"] == "ansicolors"
@@ -142,6 +146,25 @@ def test_pex_lockfile_generation(no_binary: bool, only_binary: bool) -> None:
         assert artifacts == [wheel]
 
 
+def test_constraints_file(rule_runner: RuleRunner) -> None:
+    rule_runner.write_files({"constraints.txt": "ansicolors==1.1.7"})
+    rule_runner.set_options(
+        [
+            "--python-resolves={'test': 'foo.lock'}",
+            "--python-resolves-to-constraints-file={'test': 'constraints.txt'}",
+        ],
+        env_inherit=PYTHON_BOOTSTRAP_ENV,
+    )
+
+    lock_entry = json.loads(
+        _generate(rule_runner=rule_runner, use_pex=True, ansicolors_version=">=1.0")
+    )
+    reqs = lock_entry["locked_resolves"][0]["locked_requirements"]
+    assert len(reqs) == 1
+    assert reqs[0]["project_name"] == "ansicolors"
+    assert reqs[0]["version"] == "1.1.7"
+
+
 def test_multiple_resolves() -> None:
     rule_runner = RuleRunner(
         rules=[

src/python/pants/backend/python/subsystems/setup.py

@@ -8,6 +8,7 @@
 import os
 from typing import Iterable, Iterator, Optional, cast
 
+from pants.core.goals.generate_lockfiles import UnrecognizedResolveNamesError
 from pants.option.option_types import (
     BoolOption,
     DictOption,
@@ -18,7 +19,7 @@
 )
 from pants.option.subsystem import Subsystem
 from pants.util.docutil import bin_name, doc_url
-from pants.util.memo import memoized_property
+from pants.util.memo import memoized_method, memoized_property
 from pants.util.strutil import softwrap
 
 logger = logging.getLogger(__name__)
@@ -204,7 +205,32 @@ class PythonSetup(Subsystem):
             using a resolve whose interpreter constraints are set to ['==3.7.*'], then
             Pants will error explaining the incompatibility.
 
-            The keys must be defined as resolves in `[python].resolves`.
+            The keys must be defined as resolves in `[python].resolves`. To change the interpreter
+            constraints for tool lockfiles, change `[tool].interpreter_constraints`, e.g.
+            `[black].interpreter_constraints`; if the tool does not have that option, it determines
+            its interpreter constraints from your user code.
+            """
+        ),
+        advanced=True,
+    )
+    _resolves_to_constraints_file = DictOption[str](
+        help=softwrap(
+            """
+            When generating a resolve's lockfile, use a constraints file to pin the version of
+            certain requirements. This is particularly useful to pin the versions of transitive
+            dependencies of your direct requirements.
+
+            See https://pip.pypa.io/en/stable/user_guide/#constraints-files for more information on
+            the format of constraint files and how constraints are applied in Pex and pip.
+
+            Expects a dictionary of resolve names from `[python].resolves` and Python tools (e.g.
+            `black` and `pytest`) to file paths for
+            constraints files. For example,
+            `{'data-science': '3rdparty/data-science-constraints.txt'}`.
+            If a resolve is not set in the dictionary, it will not use a constraints file.
+
+            Note: Only takes effect if you use Pex lockfiles. Use the default
+            `[python].lockfile_generator = "pex"` and run the `generate-lockfiles` goal.
             """
         ),
         advanced=True,
@@ -491,21 +517,33 @@ def generate_lockfiles_with_pex(self) -> bool:
     @memoized_property
     def resolves_to_interpreter_constraints(self) -> dict[str, tuple[str, ...]]:
         result = {}
+        unrecognized_resolves = []
         for resolve, ics in self._resolves_to_interpreter_constraints.items():
             if resolve not in self.resolves:
-                raise KeyError(
-                    softwrap(
-                        f"""
-                        Unrecognized resolve name in the option
-                        `[python].resolves_to_interpreter_constraints`: {resolve}. Each
-                        key must be one of the keys in `[python].resolves`:
-                        {sorted(self.resolves.keys())}
-                        """
-                    )
-                )
+                unrecognized_resolves.append(resolve)
             result[resolve] = tuple(ics)
+        if unrecognized_resolves:
+            raise UnrecognizedResolveNamesError(
+                unrecognized_resolves,
+                self.resolves.keys(),
+                description_of_origin="the option `[python].resolves_to_interpreter_constraints`",
+            )
         return result
 
+    @memoized_method
+    def resolves_to_constraints_file(
+        self, all_tool_resolve_names: tuple[str, ...]
+    ) -> dict[str, str]:
+        all_valid_resolves = {*self.resolves, *all_tool_resolve_names}
+        unrecognized_resolves = set(self._resolves_to_constraints_file.keys()) - all_valid_resolves
+        if unrecognized_resolves:
+            raise UnrecognizedResolveNamesError(
+                sorted(unrecognized_resolves),
+                all_valid_resolves,
+                description_of_origin="the option `[python].resolves_to_constraints_file`",
+            )
+        return self._resolves_to_constraints_file
+
     def resolve_all_constraints_was_set_explicitly(self) -> bool:
         return not self.options.is_default("resolve_all_constraints")