Support Pex's --path-mapping with lockfiles for better local requir…

…ement support (#16584)

Closes #16416.

As explained in the new docs, some users will be able to use fixed paths, which is the simplest approach. But others may have different paths per machine, e.g. using `/User/<username>` in the path. `--path-mappings` allows lockfile generators & consumers to set this on a per-machine basis.

## Tip to set to a common value

If you can use a common location and Pants's interpolation, e.g. `%(buildroot)s/wheels_dir`, then the config only needs to be set once. So, we encourage this approach.

This is deemed adequate compared to some other enhancements we could make:

- Trying to auto-detect when we can derive a common value.
- Eagerly erroring if users haven't appropriately configured `[python-repos].path_mapping`.

## Global option vs. per-resolve

The design doc at https://docs.google.com/document/d/1HAvpSNvNAHreFfvTAXavZGka-A3WWvPuH0sMjGUCo48/edit proposed having all resolve-related config be configurable on a per-resolve basis, rather than globally. We implemented this idea for a few options, but decided to not implement it (yet?) for `[python-repos].{find_links,indexes}` #16530.

Given that it is only possible to have a global setting for `[python-repos].find_links`, I could not come up with a compelling reason to let you have path mappings be per-resolve. It only complicated the user interface and `help` message because users have to think about resolves.

If we do end up implementing #16530, then it would make sense to also have this option be per-resolve. Until then, there is little value.

docs/markdown/Python/python/python-third-party-dependencies.md

@@ -422,7 +422,7 @@ setuptools
 mongomock
 ```
 
-### Version control and local requirements
+### Version control requirements
 
 You can install requirements from version control using two styles:
 
@@ -435,18 +435,6 @@ You can install requirements from version control using two styles:
   - `Django@ git+https://github.com/django/django.git@stable/2.1.x`
   - `Django@ git+https://github.com/django/django.git@fd209f62f1d83233cc634443cfac5ee4328d98b8`
 
-You can also install from local files using [PEP 440 direct references](https://www.python.org/dev/peps/pep-0440/#direct-references). You must use an absolute path to the file, and you should ensure that the file exists on your machine.
-
-```
-Django @ file:///Users/pantsbuild/prebuilt_wheels/django-3.1.1-py3-none-any.whl
-```
-
-> 🚧 Local file requirements do not yet work with lockfiles
-> 
-> Pex lockfiles will soon support local file requirements.
-> 
-> In the meantime, the workaround is to host the files in a private repository / index and load it with `[python-repos]`.
-
 > 📘 Version control via SSH
 > 
 > When using version controlled direct references hosted on private repositories with SSH access:
@@ -500,7 +488,7 @@ option, you can either use:
 
 * a URL to an HTML file with links to wheel and/or sdist files, or
 * a `file://` absolute path to an HTML file with links, or to a local directory with wheel and/or
-  sdist files.
+  sdist files. See the section on local requirements below.
 
 ```toml
 [python-repos]
@@ -533,6 +521,104 @@ the user:
 indexes.add = ["http://$USERNAME:[email protected]/index"]
 ```
 
+### Local requirements
+
+There are two ways to specify local requirements from the filesystem:
+
+* [PEP 440 direct references](https://www.python.org/dev/peps/pep-0440/#direct-references)
+
+```python 3rdparty/python
+python_requirement(
+    name="django",
+    # Use an absolute path to a .whl or sdist file.
+    requirements=["Django @ file:///Users/pantsbuild/prebuilt_wheels/django-3.1.1-py3-none-any.whl"],
+)
+
+# Reminder: we could also put this requirement string in requirements.txt and use the
+# `python_requirements` target generator.
+```
+
+* The option `[python-repos].find_links`
+
+```toml pants.toml
+[python-repos]
+# Use an absolute path to a directory containing `.whl` and/or sdist files.
+find_links = ["file:///Users/pantsbuild/prebuilt_wheels"]
+```
+```shell
+❯ ls /Users/pantsbuild/prebuilt_wheels
+ansicolors-1.1.8-py2.py3-none-any.whl
+django-3.1.1-py3-none-any.whl
+```
+```python 3rdparty/BUILD
+# Use normal requirement strings, i.e. without file paths.
+python_requirement(name="ansicolors", requirements=["ansicolors==1.1.8"])
+python_requirement(name="django", requirements=["django>=3.1,<3.2"])
+
+# Reminder: we could also put these requirement strings in requirements.txt and use the
+# `python_requirements` target generator
+```
+
+Unlike PEP 440 direct references, `[python-repos].find_links` allows you to use multiple artifacts
+for the same project name. For example, you can include multiple `.whl` and sdist files for the same
+project in the directory; if `[python-repos].indexes` is still set, then Pex/pip may use
+artifacts both from indexes like PyPI and from your local `--find-links`.
+
+Both approaches require using absolute paths, and the files must exist on your machine. This is
+usually fine when locally iterating and debugging. This approach also works well if your entire
+team can use the same fixed location. Otherwise, see the below section.
+
+#### Working around absolute paths
+
+If you need to share the lockfile on different machines, and you cannot use the same
+absolute path, then you can use the option
+`[python-repos].path_mappings` along with `[python-repos].find_links`. (`path_mappings` is not
+intended for PEP 440 direct requirements.)
+
+The `path_mappings` option allows you to substitute a portion of the absolute path with a logical
+name, which can be set to a different value than your
+teammates. For example, the path
+`file:///Users/pantsbuild/prebuilt_wheels/django-3.1.1-py3-none-any.whl` could become
+`file://${WHEELS_DIR}/django-3.1.1-py3-none-any.whl`, where each Pants user defines what
+`WHEELS_DIR` should be on their machine.
+
+This feature only works when using Pex lockfiles via `[python].resolves` and for tool lockfiles
+like Pytest and Black.
+
+`[python-repos].path_mappings` expects values in the form `NAME|PATH`, e.g.
+`WHEELS_DIR|/Users/pantsbuild/prebuilt_wheels`. Also, still use an absolute path for
+`[python-repos].find_links`.
+
+If possible, we recommend using a common file location for your whole team, and leveraging [Pants's
+interpolation](doc:options#config-file-interpolation), so that you avoid each user needing to
+manually configure `[python-repos].path_mappings` and `[python-repos].find_links`.
+For example, in `pants.toml`, you could set `[python-repos].path_mappings` to 
+`WHEELS_DIR|%(buildroot)s/python_wheels` and `[python-repos].find_links` to
+`%(buildroot)s/python_wheels`. Then, as long as every user has the folder `python_wheels` in the
+root of the repository, things will work without additional configuration. Or, you could use a
+value like `%(env.HOME)s/pants_wheels` for the path `~/pants_wheels`.
+
+```toml pants.toml
+[python-repos]
+# No one needs to change these values, as long as they can use the same shared location.
+find_links = ["file://%(buildroot)s/prebuilt_wheels"]
+path_mappings = ["WHEELS_DIR|%(buildroot)s/prebuilt_wheels"]
+```
+
+If you cannot use a common file location via interpolation, then we recommend setting these options
+in a [`.pants.rc` file](doc:options#pantsrc-file). Every teammate will need to set this up for their
+machine.
+
+```toml .pants.rc
+[python-repos]
+# Each user must set both of these to the absolute paths on their machines.
+find_links = ["file:///Users/pantsbuild/prebuilt_wheels"]
+path_mappings = ["WHEELS_DIR|/Users/pantsbuild/prebuilt_wheels"]
+```
+
+After initially setting up `[python-repos].path_mappings` and `[python-repos].find_links`, run
+`./pants generate-lockfiles` or `./pants generate-lockfiles --resolve=<resolve-name>`. You
+should see the `path_mappings` key set in the lockfile's JSON.
 
 ### Constraints files
 

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

@@ -169,7 +169,7 @@ async def _setup_pip_args_and_constraints_file(
 ) -> _PipArgsAndConstraintsSetup:
     resolve_config = await Get(ResolvePexConfig, ResolvePexConfigRequest(resolve_name))
 
-    args = list(resolve_config.indexes_and_find_links_and_manylinux_pex_args())
+    args = list(resolve_config.pex_args())
     digests = []
 
     # This feature only works with Pex lockfiles.

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

@@ -8,6 +8,7 @@
 from pants.base.deprecated import resolve_conflicting_options
 from pants.option.option_types import StrListOption
 from pants.option.subsystem import Subsystem
+from pants.util.docutil import doc_url
 from pants.util.strutil import softwrap
 
 
@@ -33,6 +34,8 @@ class PythonRepos(Subsystem):
             sdist files. Local paths must be absolute, and can either be to an HTML file with
             links or to a directory with `.whl` and/or sdist files, e.g.
             `file:///Users/pantsbuild/prebuilt_wheels`.
+
+            For local paths, you may want to use the option `[python-repos].path_mappings`.
             """
         )
     )
@@ -60,6 +63,42 @@ class PythonRepos(Subsystem):
         advanced=True,
     )
 
+    path_mappings = StrListOption(
+        help=softwrap(
+            f"""
+            Mappings to facilitate using local Python requirements when the absolute file paths
+            are different on different users' machines. For example, the
+            path `file:///Users/pantsbuild/prebuilt_wheels/django-3.1.1-py3-none-any.whl` could
+            become `file://${{WHEELS_DIR}}/django-3.1.1-py3-none-any.whl`, where each user can
+            configure what WHEELS_DIR points to on their machine.
+
+            Expects values in the form `NAME|PATH`, e.g.
+            `WHEELS_DIR|/Users/pantsbuild/prebuilt_wheels`. You can specify multiple
+            entries in the list.
+
+            This feature is intended to be used with `[python-repos].find_links`, rather than PEP
+            440 direct reference requirements (see
+            {doc_url("python-third-party-dependencies#local-requirements")}.
+            `[python-repos].find_links` must be configured to a valid absolute path for the
+            current machine.
+
+            Tip: you can avoid each user needing to manually configure this option and
+            `[python-repos].find_links` by using a common file location, along with Pants's
+            interpolation support ({doc_url('options#config-file-interpolation')}. For example,
+            in `pants.toml`, you could set both options to `%(buildroot)s/python_wheels`
+            to point to the directory `python_wheels` in the root of
+            your repository; or, use the path `%(env.HOME)s/pants_wheels` for the path
+            `~/pants_wheels`. If you are not able to use a common path like this, then we
+            recommend setting that each user set these options via a `.pants.rc` file
+            ({doc_url('options#pantsrc-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,
+    )
+
     @property
     def find_links(self) -> tuple[str, ...]:
         return cast(

src/python/pants/backend/python/util_rules/BUILD

@@ -8,6 +8,6 @@ python_tests(
     overrides={
         "local_dists_test.py": {"timeout": 120},
         "pex_from_targets_test.py": {"timeout": 200},
-        "pex_test.py": {"timeout": 330},
+        "pex_test.py": {"timeout": 400},
     },
 )

src/python/pants/backend/python/util_rules/pex.py

@@ -416,12 +416,8 @@ async def _setup_pex_requirements(
         resolve_name = None
     resolve_config = await Get(ResolvePexConfig, ResolvePexConfigRequest(resolve_name))
 
-    pex_lock_resolver_args = list(resolve_config.indexes_and_find_links_and_manylinux_pex_args())
-    pip_resolver_args = [
-        *resolve_config.indexes_and_find_links_and_manylinux_pex_args(),
-        "--resolver-version",
-        "pip-2020-resolver",
-    ]
+    pex_lock_resolver_args = list(resolve_config.pex_args())
+    pip_resolver_args = [*resolve_config.pex_args(), "--resolver-version", "pip-2020-resolver"]
 
     if isinstance(request.requirements, EntireLockfile):
         lockfile = await Get(LoadedLockfile, LoadedLockfileRequest(request.requirements.lockfile))

src/python/pants/backend/python/util_rules/pex_requirements.py

@@ -304,8 +304,14 @@ class ResolvePexConfig:
     constraints_file: ResolvePexConstraintsFile | None
     only_binary: FrozenOrderedSet[PipRequirement]
     no_binary: FrozenOrderedSet[PipRequirement]
+    path_mappings: tuple[str, ...]
 
-    def indexes_and_find_links_and_manylinux_pex_args(self) -> Iterator[str]:
+    def pex_args(self) -> Iterator[str]:
+        """Arguments for Pex for indexes/--find-links, manylinux, and path mappings.
+
+        Does not include arguments for constraints files, --only-binary, and --no-binary, which must
+        be set up independently.
+        """
         # NB: In setting `--no-pypi`, we rely on the default value of `[python-repos].indexes`
         # including PyPI, which will override `--no-pypi` and result in using PyPI in the default
         # case. Why set `--no-pypi`, then? We need to do this so that
@@ -321,6 +327,8 @@ def indexes_and_find_links_and_manylinux_pex_args(self) -> Iterator[str]:
         else:
             yield "--no-manylinux"
 
+        yield from (f"--path-mapping={v}" for v in self.path_mappings)
+
 
 @dataclass(frozen=True)
 class ResolvePexConfigRequest(EngineAwareParameter):
@@ -352,6 +360,7 @@ async def determine_resolve_pex_config(
             constraints_file=None,
             no_binary=FrozenOrderedSet(),
             only_binary=FrozenOrderedSet(),
+            path_mappings=python_repos.path_mappings,
         )
 
     all_python_tool_resolve_names = tuple(
@@ -418,6 +427,7 @@ async def determine_resolve_pex_config(
         constraints_file=constraints_file,
         no_binary=FrozenOrderedSet(no_binary),
         only_binary=FrozenOrderedSet(only_binary),
+        path_mappings=python_repos.path_mappings,
     )
 
 

src/python/pants/backend/python/util_rules/pex_requirements_test.py

@@ -153,6 +153,7 @@ def test_validate_tool_lockfiles(
             only_binary=FrozenOrderedSet(
                 [PipRequirement.parse("not-bdist" if invalid_only_binary else "bdist")]
             ),
+            path_mappings=(),
         ),
     )
 
@@ -273,6 +274,7 @@ def test_validate_user_lockfiles(
             only_binary=FrozenOrderedSet(
                 [PipRequirement.parse("not-bdist" if invalid_only_binary else "bdist")]
             ),
+            path_mappings=(),
         ),
     )