Repository: pex-tool/pex

Issue:
Add ability to update a lockfile without modifying existing entries
Sub-task of pantsbuild/pants#15704

When I add a new dependency to my project, I'd like to update my PEX lockfile with the dependency's version & transitive deps _without_ updating any of the existing entries in the lock. If it's not possible to do this (i.e. because the new dep has a transitive requirement that conflicts with something already in the lock), I'd like the command to fail (perhaps unless `-p the_transitive_dep` is passed to allow upgrading that one dependency).

Without this feature, I typically have to manually revert changes to my project's lockfile to remove "unrelated" changes. This is an error-prone / easy-to-forget step, so accidental upgrades sometimes sneak through to code review - from there they either cause confusion on the reviewer's part ("why are you upgrading X when the PR only talks about adding Y?"), or they pass through review and muddle our commit history.


Discussion:

Resolving patch:
diff --git a/pex/cli/commands/lock.py b/pex/cli/commands/lock.py
--- a/pex/cli/commands/lock.py
+++ b/pex/cli/commands/lock.py
@@ -5,7 +5,7 @@
 
 import sys
 from argparse import ArgumentParser, _ActionsContainer
-from collections import OrderedDict, defaultdict
+from collections import defaultdict
 
 from pex.argparse import HandleBoolAction
 from pex.cli.command import BuildTimeCommand
@@ -13,6 +13,7 @@
 from pex.common import pluralize
 from pex.dist_metadata import Requirement, RequirementParseError
 from pex.enum import Enum
+from pex.orderedset import OrderedSet
 from pex.resolve import requirement_options, resolver_options, target_options
 from pex.resolve.locked_resolve import LockConfiguration, LockedResolve, LockStyle
 from pex.resolve.lockfile import json_codec
@@ -167,7 +168,11 @@ def _add_update_arguments(cls, update_parser):
             action="append",
             default=[],
             type=str,
-            help="Just attempt to update these projects in the lock, leaving all others unchanged.",
+            help=(
+                "Just attempt to update these projects in the lock, leaving all others unchanged. "
+                "If the projects aren't already in the lock, attempt to add them as top-level"
+                "requirements leaving all others unchanged."
+            ),
         )
         update_parser.add_argument(
             "--strict",
@@ -373,28 +378,6 @@ def _update(self):
             return Error("Failed to parse project requirement to update: {err}".format(err=e))
 
         lock_file_path, lock_file = self._load_lockfile()
-
-        if updates:
-            updates_by_project_name = OrderedDict(
-                (update.project_name, update) for update in updates
-            )
-            for locked_resolve in lock_file.locked_resolves:
-                for locked_requirement in locked_resolve.locked_requirements:
-                    updates_by_project_name.pop(locked_requirement.pin.project_name, None)
-                    if not updates_by_project_name:
-                        break
-            if updates_by_project_name:
-                return Error(
-                    "The following updates were requested but there were no matching locked "
-                    "requirements found in {lock_file}:\n{updates}".format(
-                        lock_file=lock_file_path,
-                        updates="\n".join(
-                            "+ {update}".format(update=update)
-                            for update in updates_by_project_name.values()
-                        ),
-                    )
-                )
-
         lock_updater = LockUpdater.create(
             lock_file=lock_file,
             repos_configuration=resolver_options.create_repos_configuration(self.options),
@@ -469,17 +452,29 @@ def _update(self):
             for project_name, version_update in resolve_update.updates.items():
                 if version_update:
                     performed_update = True
-                    print(
-                        "{lead_in} {project_name} from {original_version} to {updated_version} in "
-                        "lock generated by {platform}.".format(
-                            lead_in="Would update" if dry_run else "Updated",
-                            project_name=project_name,
-                            original_version=version_update.original,
-                            updated_version=version_update.updated,
-                            platform=platform,
-                        ),
-                        file=output,
-                    )
+                    if version_update.original:
+                        print(
+                            "{lead_in} {project_name} from {original_version} to {updated_version} "
+                            "in lock generated by {platform}.".format(
+                                lead_in="Would update" if dry_run else "Updated",
+                                project_name=project_name,
+                                original_version=version_update.original,
+                                updated_version=version_update.updated,
+                                platform=platform,
+                            ),
+                            file=output,
+                        )
+                    else:
+                        print(
+                            "{lead_in} {project_name} {updated_version} to lock generated by "
+                            "{platform}.".format(
+                                lead_in="Would add" if dry_run else "Added",
+                                project_name=project_name,
+                                updated_version=version_update.updated,
+                                platform=platform,
+                            ),
+                            file=output,
+                        )
                 else:
                     print(
                         "There {tense} no updates for {project_name} in lock generated by "
@@ -491,22 +486,37 @@ def _update(self):
                         file=output,
                     )
         if performed_update:
+            original_locked_project_names = {
+                locked_requirement.pin.project_name
+                for locked_resolve in lock_file.locked_resolves
+                for locked_requirement in locked_resolve.locked_requirements
+            }
+            new_requirements = OrderedSet(
+                update
+                for update in updates
+                if update.project_name not in original_locked_project_names
+            )
             constraints_by_project_name.update(
                 (constraint.project_name, constraint) for constraint in updates
             )
-
-        if performed_update and not dry_run:
-            with open(lock_file_path, "w") as fp:
-                self._dump_lockfile(
-                    lock_file=attr.evolve(
-                        lock_file,
-                        pex_version=__version__,
-                        constraints=SortedTuple(constraints_by_project_name.values()),
-                        locked_resolves=SortedTuple(
-                            resolve_update.updated_resolve
-                            for resolve_update in lock_update.resolves
+            for requirement in new_requirements:
+                constraints_by_project_name.pop(requirement.project_name, None)
+            requirements = OrderedSet(lock_file.requirements)
+            requirements.update(new_requirements)
+
+            if not dry_run:
+                with open(lock_file_path, "w") as fp:
+                    self._dump_lockfile(
+                        lock_file=attr.evolve(
+                            lock_file,
+                            pex_version=__version__,
+                            requirements=SortedTuple(requirements, key=str),
+                            constraints=SortedTuple(constraints_by_project_name.values(), key=str),
+                            locked_resolves=SortedTuple(
+                                resolve_update.updated_resolve
+                                for resolve_update in lock_update.resolves
+                            ),
                         ),
-                    ),
-                    output=fp,
-                )
+                        output=fp,
+                    )
         return Ok()
diff --git a/pex/resolve/lockfile/updater.py b/pex/resolve/lockfile/updater.py
--- a/pex/resolve/lockfile/updater.py
+++ b/pex/resolve/lockfile/updater.py
@@ -11,6 +11,7 @@
 from pex.common import pluralize
 from pex.dist_metadata import Requirement
 from pex.network_configuration import NetworkConfiguration
+from pex.orderedset import OrderedSet
 from pex.pep_440 import Version
 from pex.pep_503 import ProjectName
 from pex.resolve.locked_resolve import LockConfiguration, LockedRequirement, LockedResolve
@@ -38,7 +39,7 @@
 
 @attr.s(frozen=True)
 class VersionUpdate(object):
-    original = attr.ib()  # type: Version
+    original = attr.ib()  # type: Optional[Version]
     updated = attr.ib()  # type: Version
 
 
@@ -82,7 +83,9 @@ def create(
             constraint.project_name: constraint for constraint in constraints
         }  # type: Mapping[ProjectName, Requirement]
 
-        update_constraints_by_project_name = {}  # type: Dict[ProjectName, Requirement]
+        update_constraints_by_project_name = (
+            OrderedDict()
+        )  # type: OrderedDict[ProjectName, Requirement]
         for update in updates:
             project_name = update.project_name
             original_constraint = original_constraints.get(project_name)
@@ -107,28 +110,36 @@ def create(
     pip_configuration = attr.ib()  # type: PipConfiguration
 
     @contextmanager
-    def _calculate_update_constraints(self, locked_resolve):
-        # type: (LockedResolve) -> Iterator[Optional[Iterable[str]]]
+    def _calculate_requirement_configuration(self, locked_resolve):
+        # type: (LockedResolve) -> Iterator[RequirementConfiguration]
         if not self.update_constraints_by_project_name:
-            yield None
+            yield RequirementConfiguration(requirements=self.original_requirements)
             return
 
+        requirements = OrderedSet(self.original_requirements)
         constraints = []
+        update_constraints_by_project_name = OrderedDict(self.update_constraints_by_project_name)
         for locked_requirement in locked_resolve.locked_requirements:
             pin = locked_requirement.pin
-            constraint = self.update_constraints_by_project_name.get(
+            constraint = update_constraints_by_project_name.pop(
                 pin.project_name, pin.as_requirement()
             )
             constraints.append(str(constraint))
+
+        # Any update constraints remaining are new requirements to resolve.
+        requirements.update(str(req) for req in update_constraints_by_project_name.values())
+
         if not constraints:
-            yield None
+            yield RequirementConfiguration(requirements=requirements)
             return
 
         with named_temporary_file(prefix="lock_update.", suffix=".constraints.txt", mode="w") as fp:
             fp.write(os.linesep.join(constraints))
             fp.flush()
             try:
-                yield [fp.name]
+                yield RequirementConfiguration(
+                    requirements=requirements, constraint_files=[fp.name]
+                )
             except ResultError as e:
                 logger.error(
                     "The following lock update constraints could not be satisfied:\n"
@@ -143,14 +154,11 @@ def update_resolve(
     ):
         # type: (...) -> Union[ResolveUpdate, Error]
 
-        with self._calculate_update_constraints(locked_resolve) as constraints_files:
+        with self._calculate_requirement_configuration(locked_resolve) as requirement_configuration:
             updated_lock_file = try_(
                 create(
                     lock_configuration=self.lock_configuration,
-                    requirement_configuration=RequirementConfiguration(
-                        requirements=self.original_requirements,
-                        constraint_files=constraints_files,
-                    ),
+                    requirement_configuration=requirement_configuration,
                     targets=targets,
                     pip_configuration=self.pip_configuration,
                 )
@@ -168,7 +176,7 @@ def update_resolve(
         for locked_requirement in locked_resolve.locked_requirements:
             original_pin = locked_requirement.pin
             project_name = original_pin.project_name
-            updated_requirement = updated_requirements_by_project_name.get(project_name)
+            updated_requirement = updated_requirements_by_project_name.pop(project_name, None)
             if not updated_requirement:
                 continue
             updated_pin = updated_requirement.pin
@@ -189,10 +197,16 @@ def update_resolve(
             elif project_name in self.update_constraints_by_project_name:
                 updates[project_name] = None
 
+        # Anything left was an addition.
+        updates.update(
+            (project_name, VersionUpdate(original=None, updated=locked_requirement.pin.version))
+            for project_name, locked_requirement in updated_requirements_by_project_name.items()
+        )
+
         return ResolveUpdate(
             updated_resolve=attr.evolve(
                 locked_resolve,
-                locked_requirements=SortedTuple(updated_requirements_by_project_name.values()),
+                locked_requirements=SortedTuple(updated_resolve.locked_requirements),
             ),
             updates=updates,
         )


Test patch:
diff --git a/tests/integration/cli/commands/test_lock.py b/tests/integration/cli/commands/test_lock.py
--- a/tests/integration/cli/commands/test_lock.py
+++ b/tests/integration/cli/commands/test_lock.py
@@ -4,6 +4,7 @@
 import json
 import os
 import re
+import shutil
 import subprocess
 import sys
 from textwrap import dedent
@@ -31,8 +32,10 @@
     PY310,
     PY_VER,
     IntegResults,
+    built_wheel,
     ensure_python_interpreter,
     make_env,
+    make_project,
     run_pex_command,
 )
 from pex.typing import TYPE_CHECKING
@@ -440,16 +443,99 @@ def test_update_noop_dry_run(lock_file_path):
     assert not result.error
 
 
-def test_update_targeted_upgrade_miss(lock_file_path):
+def test_update_targeted_add(lock_file_path):
     # type: (str) -> None
-    result = run_lock_update_for_py310("-p", "not_in_lock==1.0", lock_file_path)
-    result.assert_failure()
-    assert not result.output
+    result = run_lock_update_for_py310("-p", "ansicolors==1.1.8", lock_file_path)
+    result.assert_success()
+
+    lock_file = json_codec.load(lock_file_path)
+    assert UPDATE_LOCKFILE.constraints == lock_file.constraints
+
+    # Since the ansicolors was a new requirement and not an update of an existing project in the
+    # lock, the lock's requirements should have been updated to include it.
+    expected_requirements = SortedTuple(
+        list(UPDATE_LOCKFILE.requirements) + [Requirement.parse("ansicolors==1.1.8")], key=str
+    )  # type: SortedTuple[Requirement]
+    assert expected_requirements == lock_file.requirements
+
+    # We expect one new locked requirement for ansicolors.
+    assert 1 == len(lock_file.locked_resolves)
+    locked_resolve = lock_file.locked_resolves[0]
+    assert len(UPDATE_LOCKFILE.locked_resolves[0].locked_requirements) + 1 == len(
+        locked_resolve.locked_requirements
+    )
+
+
+def test_update_targeted_add_transitive(
+    tmpdir,  # type: Any
+    lock_file_path,  # type: str
+):
+    # type: (...) -> None
+
+    find_links = os.path.join(str(tmpdir), "find_links")
+    os.mkdir(find_links)
+
+    with built_wheel(
+        name="adds_two_new_projects_but_just_one_top_level_req",
+        version="1.2.3",
+        install_reqs=["ansicolors==1.1.8"],
+        universal=True,
+    ) as adds_two_new_projects_but_just_one_top_level_req:
+        shutil.move(
+            adds_two_new_projects_but_just_one_top_level_req,
+            os.path.join(
+                find_links, os.path.basename(adds_two_new_projects_but_just_one_top_level_req)
+            ),
+        )
+    result = run_lock_update_for_py310(
+        "-f",
+        find_links,
+        "-p",
+        "adds_two_new_projects_but_just_one_top_level_req>1",
+        lock_file_path,
+        "--indent",
+        "2",
+    )
+    result.assert_success()
+
+    lock_file = json_codec.load(lock_file_path)
+    assert UPDATE_LOCKFILE.constraints == lock_file.constraints
+
+    # Since the adds_two_new_projects_but_just_one_top_level_req was a new top-level (directly
+    # requested) requirement and not an update of an existing project in the lock, the lock's
+    # requirements should have been updated to include it.
+    expected_requirements = SortedTuple(
+        list(UPDATE_LOCKFILE.requirements)
+        + [Requirement.parse("adds_two_new_projects_but_just_one_top_level_req>1")],
+        key=str,
+    )  # type: SortedTuple[Requirement]
+    assert expected_requirements == lock_file.requirements
+
+    # We expect two new locked requirements for adds_two_new_projects_but_just_one_top_level_req.
+    assert 1 == len(lock_file.locked_resolves)
+    locked_resolve = lock_file.locked_resolves[0]
+    assert len(UPDATE_LOCKFILE.locked_resolves[0].locked_requirements) + 2 == len(
+        locked_resolve.locked_requirements
+    )
+    assert {
+        Pin(ProjectName("adds_two_new_projects_but_just_one_top_level_req"), Version("1.2.3")),
+        Pin(ProjectName("ansicolors"), Version("1.1.8")),
+    }.issubset(
+        {locked_requirement.pin for locked_requirement in locked_resolve.locked_requirements}
+    )
+
+
+def test_update_targeted_add_dry_run(lock_file_path):
+    # type: (str) -> None
+    result = run_lock_update_for_py310("-n", "-p", "ansicolors==1.1.8", lock_file_path)
+    result.assert_success()
     assert (
-        "The following updates were requested but there were no matching locked requirements "
-        "found in {lock_file}:\n"
-        "+ not_in_lock==1.0\n".format(lock_file=lock_file_path)
-    ) == result.error
+        "Would add ansicolors 1.1.8 to lock generated by cp38-cp38-manylinux_2_33_x86_64.\n"
+    ) == result.output
+    assert not result.error
+    assert UPDATE_LOCKFILE == json_codec.load(
+        lock_file_path
+    ), "A dry run update should not have updated the lock file."
 
 
 def test_update_targeted_upgrade(lock_file_path):
@@ -494,6 +580,23 @@ def test_update_targeted_upgrade_dry_run(lock_file_path):
     ), "A dry run update should not have updated the lock file."
 
 
+def test_update_targeted_mixed_dry_run(lock_file_path):
+    # type: (str) -> None
+    result = run_lock_update_for_py310(
+        "-n", "-p", "urllib3<1.26.7", "-p", "ansicolors==1.1.8", lock_file_path
+    )
+    result.assert_success()
+    assert (
+        "Would update urllib3 from 1.25.11 to 1.26.6 in lock generated by "
+        "cp38-cp38-manylinux_2_33_x86_64.\n"
+        "Would add ansicolors 1.1.8 to lock generated by cp38-cp38-manylinux_2_33_x86_64.\n"
+    ) == result.output
+    assert not result.error
+    assert UPDATE_LOCKFILE == json_codec.load(
+        lock_file_path
+    ), "A dry run update should not have updated the lock file."
+
+
 def test_update_targeted_downgrade(lock_file_path):
     # type: (str) -> None
     result = run_lock_update_for_py310("-p", "urllib3<1.25", lock_file_path)
@@ -580,6 +683,74 @@ def test_update_targeted_impossible(
     )
 
 
+def test_update_add_impossible(
+    lock_file_path,  # type: str
+    tmpdir,  # type: Any
+    py310,
+):
+    # type: (...) -> None
+
+    find_links = os.path.join(str(tmpdir), "find_links")
+    os.mkdir(find_links)
+
+    with built_wheel(
+        name="conflicting_certifi_requirement",
+        version="1.2.3",
+        install_reqs=["certifi<2017.4.17"],
+        universal=True,
+    ) as conflicting_certifi_requirement:
+        shutil.move(
+            conflicting_certifi_requirement,
+            os.path.join(find_links, os.path.basename(conflicting_certifi_requirement)),
+        )
+
+    result = run_lock_update_for_py310(
+        "-f", find_links, "-p", "conflicting_certifi_requirement", lock_file_path
+    )
+    result.assert_failure()
+    assert not result.output
+
+    error_lines = result.error.splitlines()
+    assert [
+        "ERROR: Could not find a version that satisfies the requirement certifi<2017.4.17 "
+        "(from conflicting-certifi-requirement)",
+        "ERROR: No matching distribution found for certifi<2017.4.17",
+        "ERROR: The following lock update constraints could not be satisfied:",
+        "certifi==2021.5.30",
+        "charset-normalizer==2.0.6",
+        "idna==3.2",
+        "requests==2.26",
+        "urllib3==1.25.11",
+        "Encountered 1 error updating {lock_file_path}:".format(lock_file_path=lock_file_path),
+    ] == error_lines[:9]
+    assert re.match(
+        r"^1\.\) {platform}: pid [\d]+ -> ".format(
+            platform=LocalInterpreter.create(PythonInterpreter.from_binary(py310)).platform.tag
+        ),
+        error_lines[9],
+    )
+
+    # The pip legacy resolver, though is not strict and will let us get away with this.
+    updated_lock_file_path = os.path.join(str(tmpdir), "lock.updated")
+    with safe_open(updated_lock_file_path, "w") as fp:
+        json.dump(
+            json_codec.as_json_data(
+                attr.evolve(UPDATE_LOCKFILE, resolver_version=ResolverVersion.PIP_LEGACY)
+            ),
+            fp,
+            sort_keys=True,
+        )
+    result = run_lock_update_for_py310(
+        "-f", find_links, "-p", "conflicting_certifi_requirement", updated_lock_file_path
+    )
+    result.assert_success()
+    assert not result.output
+    assert (
+        "Added conflicting-certifi-requirement 1.2.3 to lock generated by "
+        "cp38-cp38-manylinux_2_33_x86_64.\n" == result.error
+    )
+
+
 DUAL_UPDATE_LOCKFILE_CONTENTS = """\
 {
   "allow_builds": true,