Added support for discovering and patching pyright config files (pa…

…ntsbuild#17771)

- Grab the the pyrightconfig.json or pyproject.toml from the project root
- Patching the config with `venv` and `extraPaths` info - before passing them into our Process.
- Pass source roots to `extraPaths`

src/python/pants/backend/python/typecheck/pyright/rules.py

@@ -3,9 +3,13 @@
 
 from __future__ import annotations
 
-from dataclasses import dataclass
+import json
+import logging
+from dataclasses import dataclass, replace
 from typing import Iterable
 
+import toml
+
 from pants.backend.javascript.subsystems.nodejs import NpxProcess
 from pants.backend.python.subsystems.setup import PythonSetup
 from pants.backend.python.target_types import (
@@ -27,19 +31,24 @@
     PythonSourceFilesRequest,
 )
 from pants.core.goals.check import CheckRequest, CheckResult, CheckResults
+from pants.core.util_rules import config_files
+from pants.core.util_rules.config_files import ConfigFiles, ConfigFilesRequest
 from pants.core.util_rules.source_files import SourceFiles, SourceFilesRequest
 from pants.engine.collection import Collection
-from pants.engine.fs import CreateDigest, FileContent
+from pants.engine.fs import CreateDigest, DigestContents, FileContent
 from pants.engine.internals.native_engine import Digest, MergeDigests
 from pants.engine.internals.selectors import MultiGet
 from pants.engine.process import FallibleProcessResult, Process
-from pants.engine.rules import Get, Rule, collect_rules, rule
+from pants.engine.rules import Get, Rule, collect_rules, rule, rule_helper
 from pants.engine.target import CoarsenedTargets, CoarsenedTargetsRequest, FieldSet
 from pants.engine.unions import UnionRule
+from pants.source.source_root import SourceRootsRequest, SourceRootsResult
 from pants.util.logging import LogLevel
 from pants.util.ordered_set import FrozenOrderedSet, OrderedSet
 from pants.util.strutil import pluralize
 
+logger = logging.getLogger(__name__)
+
 
 @dataclass(frozen=True)
 class PyrightFieldSet(FieldSet):
@@ -71,6 +80,62 @@ class PyrightPartitions(Collection[PyrightPartition]):
     pass
 
 
+@rule_helper
+async def _patch_config_file(
+    config_files: ConfigFiles, venv_dir: str, source_roots: Iterable[str]
+) -> Digest:
+    """Patch the Pyright config file to use the incoming venv directory (from
+    requirements_venv_pex). If there is no config file, create a dummy pyrightconfig.json with the
+    `venv` key populated.
+
+    The incoming venv directory works alongside the `--venv-path` CLI argument.
+
+    Additionally, add source roots to the `extraPaths` key in the config file.
+    """
+
+    source_roots_list = list(source_roots)
+    if not config_files.snapshot.files:
+        # venv workaround as per: https://github.com/microsoft/pyright/issues/4051
+        generated_config = {"venv": venv_dir, "extraPaths": source_roots_list}
+        return await Get(
+            Digest,
+            CreateDigest(
+                [
+                    FileContent(
+                        "pyrightconfig.json",
+                        json.dumps(generated_config).encode(),
+                    )
+                ]
+            ),
+        )
+
+    config_contents = await Get(DigestContents, Digest, config_files.snapshot.digest)
+    new_files: list[FileContent] = []
+    for file in config_contents:
+        # This only supports a single json config file in the root of the project
+        # https://github.com/pantsbuild/pants/issues/17816 tracks supporting multiple config files and workspaces
+        if file.path == "pyrightconfig.json":
+            json_config = json.loads(file.content)
+            json_config["venv"] = venv_dir
+            json_extra_paths: list[str] = json_config.get("extraPaths", [])
+            json_config["extraPaths"] = list(OrderedSet(json_extra_paths + source_roots_list))
+            new_content = json.dumps(json_config).encode()
+            new_files.append(replace(file, content=new_content))
+
+        # This only supports a single pyproject.toml file in the root of the project
+        # https://github.com/pantsbuild/pants/issues/17816 tracks supporting multiple config files and workspaces
+        elif file.path == "pyproject.toml":
+            toml_config = toml.loads(file.content.decode())
+            pyright_config = toml_config["tool"]["pyright"]
+            pyright_config["venv"] = venv_dir
+            toml_extra_paths: list[str] = pyright_config.get("extraPaths", [])
+            pyright_config["extraPaths"] = list(OrderedSet(toml_extra_paths + source_roots_list))
+            new_content = toml.dumps(toml_config).encode()
+            new_files.append(replace(file, content=new_content))
+
+    return await Get(Digest, CreateDigest(new_files))
+
+
 @rule(
     desc="Pyright typecheck each partition based on its interpreter_constraints",
     level=LogLevel.DEBUG,
@@ -81,26 +146,40 @@ async def pyright_typecheck_partition(
     pex_environment: PexEnvironment,
 ) -> CheckResult:
 
-    root_sources = await Get(
+    root_sources_get = Get(
         SourceFiles,
         SourceFilesRequest(fs.sources for fs in partition.field_sets),
     )
 
     # Grab the inferred and supporting files for the root source files to be typechecked
-    coarsened_sources = await Get(
+    coarsened_sources_get = Get(
         PythonSourceFiles, PythonSourceFilesRequest(partition.root_targets.closure())
     )
 
     # See `requirements_venv_pex` for how this will get wrapped in a `VenvPex`.
-    requirements_pex = await Get(
+    requirements_pex_get = Get(
         Pex,
         RequirementsPexRequest(
             (fs.address for fs in partition.field_sets),
             hardcoded_interpreter_constraints=partition.interpreter_constraints,
         ),
     )
 
-    requirements_venv_pex = await Get(
+    # Look for any/all of the Pyright configuration files (the config is modified below for the `venv` workaround)
+    config_files_get = Get(
+        ConfigFiles,
+        ConfigFilesRequest,
+        pyright.config_request(),
+    )
+
+    root_sources, coarsened_sources, requirements_pex, config_files = await MultiGet(
+        root_sources_get,
+        coarsened_sources_get,
+        requirements_pex_get,
+        config_files_get,
+    )
+
+    requirements_venv_pex_get = Get(
         VenvPex,
         PexRequest(
             output_filename="requirements_venv.pex",
@@ -110,17 +189,24 @@ async def pyright_typecheck_partition(
         ),
     )
 
-    # venv workaround as per: https://github.com/microsoft/pyright/issues/4051
-    dummy_config_digest = await Get(
-        Digest,
-        CreateDigest(
-            [
-                FileContent(
-                    "pyrightconfig.json",
-                    f'{{ "venv": "{requirements_venv_pex.venv_rel_dir}" }}'.encode(),
-                )
-            ]
-        ),
+    source_roots_get = Get(
+        SourceRootsResult,
+        SourceRootsRequest,
+        SourceRootsRequest.for_files(root_sources.snapshot.files),
+    )
+
+    requirements_venv_pex, source_roots = await MultiGet(
+        requirements_venv_pex_get,
+        source_roots_get,
+    )
+
+    # Patch the config file to use the venv directory from the requirements pex,
+    # and add source roots to the `extraPaths` key in the config file.
+    unique_source_roots = FrozenOrderedSet(
+        [root.path for root in source_roots.path_to_root.values()]
+    )
+    patched_config_digest = await _patch_config_file(
+        config_files, requirements_venv_pex.venv_rel_dir, unique_source_roots
     )
 
     input_digest = await Get(
@@ -129,7 +215,7 @@ async def pyright_typecheck_partition(
             [
                 coarsened_sources.source_files.snapshot.digest,
                 requirements_venv_pex.digest,
-                dummy_config_digest,
+                patched_config_digest,
             ]
         ),
     )
@@ -214,6 +300,7 @@ async def pyright_typecheck(
 def rules() -> Iterable[Rule | UnionRule]:
     return (
         *collect_rules(),
+        *config_files.rules(),
         *pex_from_targets.rules(),
         UnionRule(CheckRequest, PyrightRequest),
     )

src/python/pants/backend/python/typecheck/pyright/rules_integration_test.py

@@ -62,6 +62,27 @@ def add(x: int, y: int) -> int:
     result = add(2.0, 3.0)
     """
 )
+# This will fail if `reportUndefinedVariable` is enabled (default).
+UNDEFINED_VARIABLE_FILE = dedent(
+    """\
+    print(foo)
+    """
+)
+
+UNDEFINED_VARIABLE_JSON_CONFIG = dedent(
+    """\
+    {
+        "reportUndefinedVariable": false
+    }
+    """
+)
+
+UNDEFINED_VARIABLE_TOML_CONFIG = dedent(
+    """\
+    [tool.pyright]
+    reportUndefinedVariable = false
+    """
+)
 
 
 def run_pyright(
@@ -116,6 +137,59 @@ def test_multiple_targets(rule_runner: RuleRunner) -> None:
     assert result[0].report == EMPTY_DIGEST
 
 
+@pytest.mark.parametrize(
+    "config_filename,config_file,exit_code",
+    (
+        ("pyrightconfig.json", UNDEFINED_VARIABLE_JSON_CONFIG, 0),
+        ("pyproject.toml", UNDEFINED_VARIABLE_TOML_CONFIG, 0),
+        ("noconfig", "", 1),
+    ),
+)
+def test_config_file(
+    rule_runner: RuleRunner, config_filename: str, config_file: str, exit_code: int
+) -> None:
+    rule_runner.write_files(
+        {
+            f"{PACKAGE}/f.py": UNDEFINED_VARIABLE_FILE,
+            f"{PACKAGE}/BUILD": "python_sources()",
+            f"{config_filename}": config_file,
+        }
+    )
+    tgt = rule_runner.get_target(Address(PACKAGE, relative_file_path="f.py"))
+    result = run_pyright(rule_runner, [tgt])
+    assert len(result) == 1
+    assert result[0].exit_code == exit_code
+
+
+def test_additional_source_roots(rule_runner: RuleRunner) -> None:
+    LIB_1_PACKAGE = f"{PACKAGE}/lib1"
+    LIB_2_PACKAGE = f"{PACKAGE}/lib2"
+    rule_runner.write_files(
+        {
+            f"{LIB_1_PACKAGE}/core/a.py": GOOD_FILE,
+            f"{LIB_1_PACKAGE}/core/BUILD": "python_sources()",
+            f"{LIB_2_PACKAGE}/core/b.py": "from core.a import add",
+            f"{LIB_2_PACKAGE}/core/BUILD": "python_sources()",
+        }
+    )
+    tgts = [
+        rule_runner.get_target(Address(f"{LIB_1_PACKAGE}/core", relative_file_path="a.py")),
+        rule_runner.get_target(Address(f"{LIB_2_PACKAGE}/core", relative_file_path="b.py")),
+    ]
+    result = run_pyright(rule_runner, tgts)
+    assert len(result) == 1
+    assert result[0].exit_code == 1
+    assert "reportMissingImports" in result[0].stdout
+
+    result = run_pyright(
+        rule_runner,
+        tgts,
+        extra_args=[f"--source-root-patterns=['{LIB_1_PACKAGE}', '{LIB_2_PACKAGE}']"],
+    )
+    assert len(result) == 1
+    assert result[0].exit_code == 0
+
+
 def test_skip(rule_runner: RuleRunner) -> None:
     rule_runner.write_files({f"{PACKAGE}/f.py": BAD_FILE, f"{PACKAGE}/BUILD": "python_sources()"})
     tgt = rule_runner.get_target(Address(PACKAGE, relative_file_path="f.py"))

src/python/pants/backend/python/typecheck/pyright/subsystem.py

@@ -5,6 +5,7 @@
 
 from pants.backend.javascript.subsystems.npx_tool import NpxToolBase
 from pants.backend.python.util_rules.interpreter_constraints import InterpreterConstraints
+from pants.core.util_rules.config_files import ConfigFilesRequest
 from pants.option.option_types import ArgsListOption, SkipOption, StrListOption
 from pants.util.strutil import softwrap
 
@@ -37,3 +38,22 @@ def interpreter_constraints(self) -> InterpreterConstraints:
         This assumes you have set the class property `register_interpreter_constraints = True`.
         """
         return InterpreterConstraints(self._interpreter_constraints)
+
+    def config_request(self) -> ConfigFilesRequest:
+        """Pyright will look for a `pyrightconfig.json` or a `pyproject.toml` (with a
+        `[tool.pyright]` section) in the project root.
+
+        `pyrightconfig.json` takes precedence if both are present.
+        Pyright's configuration content is specified here:
+        https://github.com/microsoft/pyright/blob/main/docs/configuration.md.
+
+        In order for Pants to work with Pyright, we modify the config file before
+        putting it in the Pyright digest. Specifically, we append source roots
+        to `extraPaths` and we overwrite `venv` to point to a pex venv.
+        """
+
+        return ConfigFilesRequest(
+            discovery=True,
+            check_existence=["pyrightconfig.json"],
+            check_content={"pyproject.toml": b"[tool.pyright"},
+        )