Respect precise file arguments with V2 ./pants test (#9120)

## Problem

When running `./pants test src/python/pants/rules/core/run_test.py`, someone who has never used Pants before would expect Pants to only test `run_test.py`. However, Pants will end up testing all of `rules/core:tests` and users must specify `./pants test src/python/pants/rules/core/run_test.py --pytest-args='-k RunTest'` to get what they want.

At the same time, we must still use the metadata from the owning target in order to prepare the test, such as resolving `dependencies`.

## Solution

Preserve the `OriginSpec` and pass it down all the way to `python_test_runner.py` so that it may pass more precise arguments to Pytest. Previously, we always passed every single file belonging to the `sources` of the target. Now, when given a filesystem spec, we will pass the files explicitly specified (which are a subset of the target's `sources`).

In order to preserve the `OriginSpec`, we make several changes.

**These changes are clunky due to the lack of multiple params in `await Get` (#7490). This PR helps to guide that design, though, while adding a useful feature needed today.**

### `AddressesWithOrigins`

The test `@goal_rule` now requests `AddressesWithOrigins`, rather than `Addresses`, so that we may preserve the original specs.

Here, we add a method to go from `AddressesWithOrigins -> AddressWithOrigin`.

### `HydratedTargetWithOrigin`

This allows us to have a hydrated target without losing the `OriginSpec` information. 

We need to be able to use `HydratedTarget` because this is how we access the underlying `TargetAdaptor` necessary to get the `TestTarget` union working. The original `AddressWithOrigin` from the `@goal_rule` isn't hydrated enough.

### `TargetAdaptorWithOrigin`

Finally, we add `TargetAdaptorWithOrigin`, along with a type for each distinct `TargetAdaptor`, such as `PythonTestsAdaptorWithOrigin`.

We need so many new classes for the union mechanism to work properly. `python_test_runner.py` must have some way to signal that it only works with `python_tests` targets, which we previously did by registering `UnionRule(TestTarget, PythonTestsAdaptor)`. But, we still need to preserve the `OriginSpec`, so we introduce `PythonTestsAdaptorWithOrigin`.

#### Rejected alternative: field on `TargetAdaptor`

One alternative design is to insert the `OriginSpec` proactively on the `TargetAdaptor`. This means that we would not need to introduce any `TargetAdaptorWithOrigin` types.

However, this approach suffers from not being typed. Given the type-driven API of rules, we want some structured way to specify that `python_test_runner.py` _must_ have the `OriginSpec` preserved, rather than checking for some optional field defined on the `TargetAdaptor`.

## Result

`./pants test src/python/pants/rules/core/run_test.py` only runs over `run_test.py` now. 

In a followup, we will add precise file args to `fmt2` and `lint2`.

src/python/pants/backend/python/rules/python_test_runner.py

@@ -3,6 +3,7 @@
 
 import os
 from dataclasses import dataclass
+from pathlib import PurePath
 from textwrap import dedent
 from typing import Optional, Tuple
 
@@ -11,18 +12,26 @@
 from pants.backend.python.rules.prepare_chrooted_python_sources import ChrootedPythonSources
 from pants.backend.python.subsystems.pytest import PyTest
 from pants.backend.python.subsystems.subprocess_environment import SubprocessEncodingEnvironment
-from pants.build_graph.address import Address
+from pants.base.specs import FilesystemResolvedSpec
 from pants.engine.addressable import Addresses
-from pants.engine.fs import Digest, DirectoriesToMerge, FileContent, InputFilesContent
+from pants.engine.fs import (
+  Digest,
+  DirectoriesToMerge,
+  FileContent,
+  InputFilesContent,
+  PathGlobs,
+  Snapshot,
+  SnapshotSubset,
+)
 from pants.engine.interactive_runner import InteractiveProcessRequest
 from pants.engine.isolated_process import ExecuteProcessRequest, FallibleExecuteProcessResult
 from pants.engine.legacy.graph import HydratedTargets, TransitiveHydratedTargets
-from pants.engine.legacy.structs import PythonTestsAdaptor
+from pants.engine.legacy.structs import PythonTestsAdaptor, PythonTestsAdaptorWithOrigin
 from pants.engine.rules import UnionRule, rule, subsystem_rule
 from pants.engine.selectors import Get
 from pants.option.global_options import GlobalOptions
 from pants.python.python_setup import PythonSetup
-from pants.rules.core.strip_source_roots import SourceRootStrippedSources
+from pants.rules.core.strip_source_roots import SourceRootStrippedSources, StripSourceRootsRequest
 from pants.rules.core.test import TestDebugRequest, TestOptions, TestResult, TestTarget
 
 
@@ -80,32 +89,33 @@ class TestTargetSetup:
 
 
 def get_packages_to_cover(
-  test_target: PythonTestsAdaptor,
+  target: PythonTestsAdaptor,
   source_root_stripped_file_paths: Tuple[str, ...],
 ) -> Tuple[str, ...]:
-  if hasattr(test_target, 'coverage'):
-    return tuple(sorted(set(test_target.coverage)))
+  if hasattr(target, 'coverage'):
+    return tuple(sorted(set(target.coverage)))
   return tuple(sorted({
     os.path.dirname(source_root_stripped_source_file_path).replace(os.sep, '.') # Turn file paths into package names.
     for source_root_stripped_source_file_path in source_root_stripped_file_paths
   }))
 
+
 @rule
 async def setup_pytest_for_target(
-  test_target: PythonTestsAdaptor,
+  target_with_origin: PythonTestsAdaptorWithOrigin,
   pytest: PyTest,
   test_options: TestOptions,
 ) -> TestTargetSetup:
+  adaptor = target_with_origin.adaptor
+  origin = target_with_origin.origin
   # TODO: Rather than consuming the TestOptions subsystem, the TestRunner should pass on coverage
   # configuration via #7490.
-  transitive_hydrated_targets = await Get[TransitiveHydratedTargets](
-    Addresses((test_target.address,))
-  )
+  transitive_hydrated_targets = await Get[TransitiveHydratedTargets](Addresses((adaptor.address,)))
   all_targets = transitive_hydrated_targets.closure
 
   resolved_requirements_pex = await Get[Pex](
     CreatePexFromTargetClosure(
-      addresses=Addresses((test_target.address,)),
+      addresses=Addresses((adaptor.address,)),
       output_filename='pytest-with-requirements.pex',
       entry_point="pytest:main",
       additional_requirements=pytest.get_requirement_strings(),
@@ -131,18 +141,36 @@ async def setup_pytest_for_target(
   # https://pytest.org/en/latest/goodpractices.html#test-discovery. In addition to a performance
   # optimization, this ensures that any transitive sources, such as a test project file named
   # test_fail.py, do not unintentionally end up being run as tests.
-  source_root_stripped_test_target_sources = await Get[SourceRootStrippedSources](
-    Address, test_target.address
+  # TODO: generalize this once it's we add precise file args to `fmt` and `lint`. Those will likely
+  # have a slightly different implementation, such as _not_ stripping source roots for most linters,
+  # so it's not clear yet what should be generalized.
+  test_files_snapshot = adaptor.sources.snapshot
+  if isinstance(origin, FilesystemResolvedSpec):
+    # NB: we ensure that `precise_files_specified` is a subset of the original target's `sources`.
+    # It's possible when given a glob filesystem spec that the spec will have resolved files not
+    # belonging to this target - those must be filtered out.
+    precise_files_specified = set(test_files_snapshot.files).intersection(origin.resolved_files)
+    test_files_snapshot = await Get[Snapshot](
+      SnapshotSubset(
+        directory_digest=test_files_snapshot.directory_digest,
+        globs=PathGlobs(sorted(precise_files_specified)),
+      )
+    )
+  test_files = await Get[SourceRootStrippedSources](
+    StripSourceRootsRequest(
+      test_files_snapshot,
+      # TODO: simply pass `address.spec_path` once `--source-unmatched` is removed.
+      representative_path=PurePath(adaptor.address.spec_path, "BUILD").as_posix(),
+    )
   )
+  test_file_names = test_files.snapshot.files
 
   coverage_args = []
-  test_target_sources_file_names = source_root_stripped_test_target_sources.snapshot.files
   if test_options.values.run_coverage:
     coveragerc_digest = await Get[Digest](InputFilesContent, get_coveragerc_input(DEFAULT_COVERAGE_CONFIG))
     directories_to_merge.append(coveragerc_digest)
     packages_to_cover = get_packages_to_cover(
-      test_target,
-      source_root_stripped_file_paths=test_target_sources_file_names,
+      adaptor, source_root_stripped_file_paths=test_file_names,
     )
     coverage_args = [
       '--cov-report=',  # To not generate any output. https://pytest-cov.readthedocs.io/en/latest/config.html
@@ -154,22 +182,22 @@ async def setup_pytest_for_target(
 
   timeout_seconds = calculate_timeout_seconds(
     timeouts_enabled=pytest.options.timeouts,
-    target_timeout=getattr(test_target, 'timeout', None),
+    target_timeout=getattr(adaptor, 'timeout', None),
     timeout_default=pytest.options.timeout_default,
     timeout_maximum=pytest.options.timeout_maximum,
   )
 
   return TestTargetSetup(
     requirements_pex=resolved_requirements_pex,
-    args=(*pytest.options.args, *coverage_args, *sorted(test_target_sources_file_names)),
+    args=(*pytest.options.args, *coverage_args, *sorted(test_file_names)),
     input_files_digest=merged_input_files,
     timeout_seconds=timeout_seconds,
   )
 
 
 @rule(name="Run pytest")
 async def run_python_test(
-  test_target: PythonTestsAdaptor,
+  target_with_origin: PythonTestsAdaptorWithOrigin,
   test_setup: TestTargetSetup,
   python_setup: PythonSetup,
   subprocess_encoding_environment: SubprocessEncodingEnvironment,
@@ -187,7 +215,7 @@ async def run_python_test(
     pex_args=test_setup.args,
     input_files=test_setup.input_files_digest,
     output_directories=('.coverage',) if test_options.values.run_coverage else None,
-    description=f'Run Pytest for {test_target.address.reference()}',
+    description=f'Run Pytest for {target_with_origin.adaptor.address.reference()}',
     timeout_seconds=test_setup.timeout_seconds if test_setup.timeout_seconds is not None else 9999,
     env=env,
   )
@@ -210,7 +238,7 @@ def rules():
     run_python_test,
     debug_python_test,
     setup_pytest_for_target,
-    UnionRule(TestTarget, PythonTestsAdaptor),
+    UnionRule(TestTarget, PythonTestsAdaptorWithOrigin),
     subsystem_rule(PyTest),
     subsystem_rule(PythonSetup),
   ]

src/python/pants/backend/python/rules/python_test_runner_integration_test.py

@@ -17,11 +17,11 @@
 from pants.backend.python.targets.python_library import PythonLibrary
 from pants.backend.python.targets.python_requirement_library import PythonRequirementLibrary
 from pants.backend.python.targets.python_tests import PythonTests
-from pants.build_graph.address import Address
+from pants.base.specs import FilesystemLiteralSpec, OriginSpec, SingleAddress
 from pants.build_graph.build_file_aliases import BuildFileAliases
 from pants.engine.fs import FileContent
 from pants.engine.interactive_runner import InteractiveRunner
-from pants.engine.legacy.structs import PythonTestsAdaptor
+from pants.engine.legacy.structs import PythonTestsAdaptor, PythonTestsAdaptorWithOrigin
 from pants.engine.rules import RootRule, subsystem_rule
 from pants.engine.selectors import Params
 from pants.python.python_requirement import PythonRequirement
@@ -46,7 +46,7 @@ class PythonTestRunnerIntegrationTest(TestBase):
 
   def write_file(self, file_content: FileContent) -> None:
     self.create_file(
-      relpath=str(PurePath(self.source_root, file_content.path)),
+      relpath=PurePath(self.source_root, file_content.path).as_posix(),
       contents=file_content.content.decode(),
     )
 
@@ -115,10 +115,12 @@ def rules(cls):
       *strip_source_roots.rules(),
       *subprocess_environment.rules(),
       subsystem_rule(TestOptions),
-      RootRule(PythonTestsAdaptor),
+      RootRule(PythonTestsAdaptorWithOrigin),
     )
 
-  def run_pytest(self, *, passthrough_args: Optional[str] = None) -> TestResult:
+  def run_pytest(
+    self, *, passthrough_args: Optional[str] = None, origin: Optional[OriginSpec] = None,
+  ) -> TestResult:
     args = [
       "--backend-packages2=pants.backend.python",
       # pin to lower versions so that we can run Python 2 tests
@@ -128,11 +130,19 @@ def run_pytest(self, *, passthrough_args: Optional[str] = None) -> TestResult:
     if passthrough_args:
       args.append(f"--pytest-args='{passthrough_args}'")
     options_bootstrapper = create_options_bootstrapper(args=args)
-    target = PythonTestsAdaptor(address=Address.parse(f"{self.source_root}:target"))
-    test_result = self.request_single_product(TestResult, Params(target, options_bootstrapper))
-    debug_request = self.request_single_product(
-      TestDebugRequest, Params(target, options_bootstrapper),
+    if origin is None:
+      origin = SingleAddress(directory=self.source_root, name="target")
+    # TODO: We must use the V1 target's `_sources_field.sources` field to set the TargetAdaptor's
+    # sources attribute. The adaptor will not auto-populate this field. However, it will
+    # auto-populate things like `dependencies` and this was not necessary before using
+    # PythonTestsAdaptorWithOrigin. Why is this necessary in test code?
+    v1_target = self.target(f"{self.source_root}:target")
+    adaptor = PythonTestsAdaptor(
+      address=v1_target.address.to_address(), sources=v1_target._sources_field.sources,
     )
+    params = Params(PythonTestsAdaptorWithOrigin(adaptor, origin), options_bootstrapper)
+    test_result = self.request_single_product(TestResult, params)
+    debug_request = self.request_single_product(TestDebugRequest, params)
     debug_result = InteractiveRunner(self.scheduler).run_local_interactive_process(debug_request.ipr)
     if test_result.status == Status.SUCCESS:
       assert debug_result.process_exit_code == 0
@@ -159,6 +169,14 @@ def test_mixed_sources(self) -> None:
     assert "test_good.py ." in result.stdout
     assert "test_bad.py F" in result.stdout
 
+  def test_precise_file_args(self) -> None:
+    self.create_python_test_target([self.good_source, self.bad_source])
+    file_arg = FilesystemLiteralSpec(PurePath(self.source_root, self.good_source.path).as_posix())
+    result = self.run_pytest(origin=file_arg)
+    assert result.status == Status.SUCCESS
+    assert "test_good.py ." in result.stdout
+    assert "test_bad.py F" not in result.stdout
+
   def test_absolute_import(self) -> None:
     self.create_basic_library()
     source = FileContent(

src/python/pants/base/specs.py

@@ -442,3 +442,6 @@ def provided_specs(self) -> Union[AddressSpecs, FilesystemSpecs]:
       if self.filesystem_specs.dependencies
       else self.address_specs
     )
+
+
+OriginSpec = Union[AddressSpec, FilesystemResolvedSpec]

src/python/pants/engine/addressable.py

@@ -5,38 +5,46 @@
 from collections.abc import MutableMapping, MutableSequence
 from dataclasses import dataclass
 from functools import update_wrapper
-from typing import Any, Set, Tuple, Type, Union
+from typing import Any, Sequence, Set, Tuple, Type
 
 from pants.base.exceptions import ResolveError
-from pants.base.specs import AddressSpec, FilesystemResolvedSpec
+from pants.base.specs import OriginSpec
 from pants.build_graph.address import Address, BuildFileAddress
 from pants.engine.objects import Collection, Resolvable, Serializable
 from pants.util.objects import TypeConstraintError
 
 
+def _assert_single_address(addresses: Sequence[Address]) -> None:
+  """Assert that exactly one address must be contained in the collection."""
+  if len(addresses) == 0:
+    raise ResolveError("No targets were matched.")
+  if len(addresses) > 1:
+    output = '\n  * '.join(address.spec for address in addresses)
+    raise ResolveError(
+      "Expected a single target, but was given multiple targets.\n\n"
+      f"Did you mean one of:\n  * {output}"
+    )
+
+
 class Addresses(Collection[Address]):
   def expect_single(self) -> Address:
-    """Assert that exactly one Address must be contained in the collection, then return it."""
-    if len(self.dependencies) == 0:
-      raise ResolveError("No targets were matched.")
-    if len(self.dependencies) > 1:
-      output = '\n  * '.join(address.spec for address in self.dependencies)
-      raise ResolveError(
-        "Expected a single target, but was given multiple targets.\n\n"
-        f"Did you mean one of:\n  * {output}"
-      )
+    _assert_single_address(self.dependencies)
     return self.dependencies[0]
 
 
 @dataclass(frozen=True)
 class AddressWithOrigin:
   """A BuildFileAddress along with the cmd-line spec it was generated from."""
   address: Address
-  origin: Union[AddressSpec, FilesystemResolvedSpec]
+  origin: OriginSpec
 
 
 class AddressesWithOrigins(Collection[AddressWithOrigin]):
-  pass
+  def expect_single(self) -> AddressWithOrigin:
+    _assert_single_address(
+      [address_with_origin.address for address_with_origin in self.dependencies]
+    )
+    return self.dependencies[0]
 
 
 class BuildFileAddresses(Collection[BuildFileAddress]):

src/python/pants/engine/build_files.py

@@ -321,5 +321,6 @@ def address_mapper_singleton() -> AddressMapper:
     address_origin_map,
     # Root rules representing parameters that might be provided via root subjects.
     RootRule(Address),
+    RootRule(AddressWithOrigin),
     RootRule(AddressSpecs),
   ]

src/python/pants/engine/legacy/graph.py

@@ -21,6 +21,7 @@
   FilesystemLiteralSpec,
   FilesystemResolvedGlobSpec,
   FilesystemSpecs,
+  OriginSpec,
   SingleAddress,
 )
 from pants.build_graph.address import Address, BuildFileAddress
@@ -411,6 +412,19 @@ class LegacyTransitiveHydratedTargets:
   closure: OrderedSet  # TODO: this is an OrderedSet[LegacyHydratedTarget]
 
 
+# TODO(#7490): Remove this once we have multiple params support so that rules can do something
+# like `await Get[TestResult](Params(Address(..), Origin(..)))`.
+@dataclass(frozen=True)
+class HydratedTargetWithOrigin:
+  """A wrapper around HydratedTarget that preserves the original spec used to resolve the target.
+
+  This is useful for precise file arguments, where a goal like `./pants test` runs over only the
+  specified file arguments rather than the whole target.
+  """
+  target: HydratedTarget
+  origin: OriginSpec
+
+
 @dataclass(frozen=True)
 class SourcesSnapshot:
   """Sources matched by command line specs, either directly via FilesystemSpecs or indirectly via
@@ -631,6 +645,14 @@ async def hydrate_target(hydrated_struct: HydratedStruct) -> HydratedTarget:
   )
 
 
+@rule
+async def hydrate_target_with_origin(
+  address_with_origin: AddressWithOrigin
+) -> HydratedTargetWithOrigin:
+  ht = await Get[HydratedTarget](Address, address_with_origin.address)
+  return HydratedTargetWithOrigin(target=ht, origin=address_with_origin.origin)
+
+
 def _eager_fileset_with_spec(
   spec_path: str, filespec: Filespec, snapshot: Snapshot, include_dirs: bool = False,
 ) -> EagerFilesetWithSpec:
@@ -797,6 +819,7 @@ def create_legacy_graph_tasks():
     legacy_transitive_hydrated_targets,
     hydrated_targets,
     hydrate_target,
+    hydrate_target_with_origin,
     find_owners,
     hydrate_sources,
     hydrate_bundles,