Fix Ruff import sorting when fix batches omit package markers

docs/notes/2.33.x.md

@@ -34,6 +34,8 @@ The `docker_image` target now supports capturing file and directory artifacts fr
 
 The Python Build Standalone backend ([pants.backend.python.providers.experimental.python_build_standalone](https://www.pantsbuild.org/stable/reference/subsystems/python-build-standalone-python-provider)) has release metadata current through PBS release [20260414](https://github.com/astral-sh/python-build-standalone/releases/tag/20260414).
 
+Fixed a bug in Ruff backend where ancestor `__init__.py` files were not included in `fix` partitions which changed how Ruff's importing sorting logic operated. Thie caused `pants fix` to see no need for import sorting even though `pants lint` flagged the need.
+
 #### Shell
 
 #### Javascript

src/python/pants/backend/python/lint/ruff/check/rules.py

@@ -3,7 +3,9 @@
 
 from __future__ import annotations
 
+import itertools
 from dataclasses import dataclass
+from pathlib import PurePosixPath
 from typing import Any
 
 from pants.backend.python.lint.ruff.check.skip_field import SkipRuffCheckField
@@ -18,13 +20,27 @@
 from pants.backend.python.util_rules import pex
 from pants.core.goals.fix import FixResult, FixTargetsRequest
 from pants.core.goals.lint import REPORT_DIR, LintResult, LintTargetsRequest
-from pants.core.util_rules.partitions import PartitionerType
+from pants.core.util_rules.partitions import Partition, PartitionerType, Partitions
 from pants.core.util_rules.source_files import SourceFilesRequest, determine_source_files
-from pants.engine.fs import DigestSubset, PathGlobs, RemovePrefix
-from pants.engine.intrinsics import digest_subset_to_digest, remove_prefix
+from pants.engine.fs import (
+    CreateDigest,
+    DigestSubset,
+    FileContent,
+    MergeDigests,
+    PathGlobs,
+    RemovePrefix,
+)
+from pants.engine.internals.graph import resolve_source_paths
+from pants.engine.intrinsics import (
+    create_digest,
+    digest_subset_to_digest,
+    digest_to_snapshot,
+    merge_digests,
+    remove_prefix,
+)
 from pants.engine.platform import Platform
-from pants.engine.rules import collect_rules, rule
-from pants.engine.target import FieldSet, Target
+from pants.engine.rules import collect_rules, concurrently, implicitly, rule
+from pants.engine.target import FieldSet, SourcesPathsRequest, Target
 from pants.util.logging import LogLevel
 from pants.util.meta import classproperty
 
@@ -59,7 +75,7 @@ def tool_id(cls) -> str:
 class RuffFixRequest(FixTargetsRequest):
     field_set_type = RuffCheckFieldSet
     tool_subsystem = Ruff  # type: ignore[assignment]
-    partitioner_type = PartitionerType.DEFAULT_SINGLE_PARTITION
+    partitioner_type = PartitionerType.CUSTOM
 
     # We don't need to include automatically added lint rules for this RuffFixRequest,
     # because these lint rules are already checked by RuffLintRequest.
@@ -74,10 +90,72 @@ def tool_id(cls) -> str:
         return RuffLintRequest.tool_id
 
 
+@dataclass(frozen=True)
+class RuffFixPartitionMetadata:
+    init_files: tuple[str, ...]
+
+    @property
+    def description(self) -> None:
+        return None
+
+
+def _ancestor_init_files(
+    files: tuple[str, ...], candidate_init_files: set[str] | frozenset[str]
+) -> tuple[str, ...]:
+    init_files = set[str]()
+    for file in files:
+        directory = PurePosixPath(file).parent
+        while directory != PurePosixPath("."):
+            init_file = str(directory / "__init__.py")
+            if init_file in candidate_init_files:
+                init_files.add(init_file)
+            directory = directory.parent
+    return tuple(sorted(init_files))
+
+
+@rule
+async def partition_ruff_fix(
+    request: RuffFixRequest.PartitionRequest, ruff: Ruff
+) -> Partitions[str, RuffFixPartitionMetadata]:
+    if ruff.skip:
+        return Partitions()
+
+    all_sources_paths = await concurrently(
+        resolve_source_paths(SourcesPathsRequest(field_set.source), **implicitly())
+        for field_set in request.field_sets
+    )
+    files = tuple(
+        sorted(
+            itertools.chain.from_iterable(
+                sources_paths.files for sources_paths in all_sources_paths
+            )
+        )
+    )
+    selected_init_files = {file for file in files if PurePosixPath(file).name == "__init__.py"}
+    metadata = RuffFixPartitionMetadata(_ancestor_init_files(files, selected_init_files))
+
+    return Partitions([Partition(files, metadata)])
+
+
 @rule(desc="Fix with `ruff check --fix`", level=LogLevel.DEBUG)
 async def ruff_fix(request: RuffFixRequest.Batch, ruff: Ruff, platform: Platform) -> FixResult:
+    # Ruff's isort rules use package marker files to classify imports. The `fix` goal may split
+    # editable files into smaller batches, so reconstruct selected ancestor `__init__.py` files as
+    # read-only context while still asking Ruff to edit only the current batch.
+    init_files = (
+        request.partition_metadata.init_files
+        if isinstance(request.partition_metadata, RuffFixPartitionMetadata)
+        else ()
+    )
+    missing_init_files = tuple(file for file in init_files if file not in request.snapshot.files)
+    init_digest = await create_digest(
+        CreateDigest(FileContent(file, b"") for file in missing_init_files)
+    )
+    snapshot = await digest_to_snapshot(
+        await merge_digests(MergeDigests((request.snapshot.digest, init_digest)))
+    )
     result = await run_ruff(
-        RunRuffRequest(snapshot=request.snapshot, mode=RuffMode.FIX),
+        RunRuffRequest(snapshot=snapshot, files=request.files, mode=RuffMode.FIX),
         ruff,
         platform,
     )
@@ -94,7 +172,11 @@ async def ruff_lint(
         SourceFilesRequest(field_set.source for field_set in request.elements)
     )
     result = await run_ruff(
-        RunRuffRequest(snapshot=source_files.snapshot, mode=RuffMode.LINT),
+        RunRuffRequest(
+            snapshot=source_files.snapshot,
+            files=source_files.snapshot.files,
+            mode=RuffMode.LINT,
+        ),
         ruff,
         platform,
     )

src/python/pants/backend/python/lint/ruff/common.py

@@ -22,6 +22,7 @@
 @dataclass(frozen=True)
 class RunRuffRequest:
     snapshot: Snapshot
+    files: tuple[str, ...]
     mode: RuffMode
 
 
@@ -71,12 +72,12 @@ async def run_ruff(
 
     result = await execute_process(
         Process(
-            argv=(exe_path, *initial_args, *conf_args, *ruff.args, *request.snapshot.files),
+            argv=(exe_path, *initial_args, *conf_args, *ruff.args, *request.files),
             input_digest=input_digest,
             immutable_input_digests={immutable_input_key: ruff_tool.digest},
-            output_files=request.snapshot.files,
+            output_files=request.files,
             output_directories=(REPORT_DIR,) if request.mode is RuffMode.LINT else (),
-            description=f"Run ruff {' '.join(initial_args)} on {pluralize(len(request.snapshot.files), 'file')}.",
+            description=f"Run ruff {' '.join(initial_args)} on {pluralize(len(request.files), 'file')}.",
             level=LogLevel.DEBUG,
         ),
         **implicitly(),

src/python/pants/backend/python/lint/ruff/format/rules.py

@@ -58,6 +58,7 @@ async def _run_ruff_fmt(
 ) -> FmtResult:
     run_ruff_request = RunRuffRequest(
         snapshot=request.snapshot,
+        files=request.files,
         mode=RuffMode.FORMAT,
     )
     result = await run_ruff(run_ruff_request, ruff, platform)

src/python/pants/backend/python/lint/ruff/rules_integration_test.py

@@ -28,7 +28,7 @@
 from pants.core.goals.fmt import FmtResult
 from pants.core.goals.lint import LintResult
 from pants.core.util_rules import config_files
-from pants.core.util_rules.partitions import _EmptyMetadata
+from pants.core.util_rules.partitions import Partitions, _EmptyMetadata
 from pants.core.util_rules.source_files import SourceFiles, SourceFilesRequest
 from pants.engine.addresses import Address
 from pants.engine.fs import DigestContents
@@ -39,6 +39,15 @@
 GOOD_FILE = 'a = "string without any placeholders"\n'
 BAD_FILE = 'a = f"string without any placeholders"\n'
 UNFORMATTED_FILE = 'a ="string without any placeholders"\n'
+IMPORTS_SHADOWED_BY_LOCAL_PACKAGE = (
+    "from opentelemetry import trace\n"
+    "\n"
+    "from pants.backend.observability.opentelemetry.config import Config\n"
+)
+IMPORTS_FIXED_FOR_LOCAL_PACKAGE = (
+    "from opentelemetry import trace\n"
+    "from pants.backend.observability.opentelemetry.config import Config\n"
+)
 
 
 @pytest.fixture
@@ -56,6 +65,7 @@ def rule_runner() -> RuleRunner:
             QueryRule(FixResult, [RuffFixRequest.Batch]),
             QueryRule(LintResult, [RuffLintRequest.Batch]),
             QueryRule(FmtResult, [RuffFormatRequest.Batch]),
+            QueryRule(Partitions, [RuffFixRequest.PartitionRequest]),
             QueryRule(SourceFiles, (SourceFilesRequest,)),
         ],
         target_types=[PythonSourcesGeneratorTarget],
@@ -88,7 +98,7 @@ def run_ruff(
         [
             RuffFixRequest.Batch(
                 "",
-                tuple(check_field_sets),
+                check_input_sources.snapshot.files,
                 partition_metadata=_EmptyMetadata(),
                 snapshot=check_input_sources.snapshot,
             ),
@@ -109,7 +119,7 @@ def run_ruff(
         [
             RuffFormatRequest.Batch(
                 "",
-                tuple(format_field_sets),
+                format_input_sources.snapshot.files,
                 partition_metadata=_EmptyMetadata(),
                 snapshot=format_input_sources.snapshot,
             )
@@ -180,6 +190,98 @@ def test_multiple_targets(rule_runner: RuleRunner) -> None:
     assert fmt_result.did_change is True
 
 
+def test_fix_preserves_init_py_package_context_for_import_sorting(rule_runner: RuleRunner) -> None:
+    rule_runner.write_files(
+        {
+            "pyproject.toml": "\n".join(
+                [
+                    "[tool.ruff.lint]",
+                    'select = ["I"]',
+                    "",
+                    "[tool.ruff.lint.isort]",
+                    'known-first-party = ["pants"]',
+                    "",
+                ]
+            ),
+            "src/python/pants/backend/observability/opentelemetry/BUILD": "python_sources()",
+            "src/python/pants/backend/observability/opentelemetry/__init__.py": "",
+            "src/python/pants/backend/observability/opentelemetry/processor.py": (
+                IMPORTS_SHADOWED_BY_LOCAL_PACKAGE
+            ),
+        }
+    )
+    rule_runner.set_options(
+        ["--backend-packages=pants.backend.python.lint.ruff"],
+        env_inherit={"PATH", "PYENV_ROOT", "HOME"},
+    )
+
+    init_tgt = rule_runner.get_target(
+        Address(
+            "src/python/pants/backend/observability/opentelemetry",
+            relative_file_path="__init__.py",
+        )
+    )
+    processor_tgt = rule_runner.get_target(
+        Address(
+            "src/python/pants/backend/observability/opentelemetry",
+            relative_file_path="processor.py",
+        )
+    )
+    init_field_set = RuffCheckFieldSet.create(init_tgt)
+    processor_field_set = RuffCheckFieldSet.create(processor_tgt)
+
+    lint_result = rule_runner.request(
+        LintResult,
+        [
+            RuffLintRequest.Batch(
+                "",
+                (init_field_set, processor_field_set),
+                partition_metadata=_EmptyMetadata(),
+            ),
+        ],
+    )
+    assert lint_result.exit_code == 1
+    assert "I001" in lint_result.stdout
+
+    partitions = rule_runner.request(
+        Partitions,
+        [RuffFixRequest.PartitionRequest((init_field_set, processor_field_set))],
+    )
+    assert len(partitions) == 1
+    partition = partitions[0]
+    assert partition.elements == (
+        "src/python/pants/backend/observability/opentelemetry/__init__.py",
+        "src/python/pants/backend/observability/opentelemetry/processor.py",
+    )
+    assert partition.metadata.init_files == (
+        "src/python/pants/backend/observability/opentelemetry/__init__.py",
+    )
+
+    fix_input_sources = rule_runner.request(
+        SourceFiles, [SourceFilesRequest([processor_field_set.source])]
+    )
+    fix_result = rule_runner.request(
+        FixResult,
+        [
+            RuffFixRequest.Batch(
+                "",
+                ("src/python/pants/backend/observability/opentelemetry/processor.py",),
+                partition_metadata=partition.metadata,
+                snapshot=fix_input_sources.snapshot,
+            ),
+        ],
+    )
+
+    assert fix_result.stdout == "Found 1 error (1 fixed, 0 remaining).\n"
+    assert fix_result.output == rule_runner.make_snapshot(
+        {
+            "src/python/pants/backend/observability/opentelemetry/processor.py": (
+                IMPORTS_FIXED_FOR_LOCAL_PACKAGE
+            )
+        }
+    )
+
+
 def test_skip_field(rule_runner: RuleRunner) -> None:
     rule_runner.write_files(
         {