Fix MyPy to not directly run over dependencies

[Eric-Arellano]

Right now, we instruct MyPy to run over not only the specified files, but all transitive deps. A user found that this results in type checking more than the user actually intended. It is true that we need all transitive deps to be in the chroot, but we shouldn't be telling MyPy to run directly on those.

FYI, how it works now

We partition MyPy based on interpreter constraints, which allows for you to run over a codebase that might have some parts only Py2 and another part only Py3:

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

Lines 60 to 65 in 9954e08

	@dataclass(frozen=True)
	class MyPyPartition:
	field_set_addresses: FrozenOrderedSet[Address]
	closure: FrozenOrderedSet[Target]
	interpreter_constraints: PexInterpreterConstraints
	python_version_already_configured: bool

All the interesting work happens on partitions, in this rule:

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

Lines 184 to 185 in 9954e08

	@rule
	async def mypy_typecheck_partition(partition: MyPyPartition, mypy: MyPy) -> TypecheckResult:

We get all the source files for the partition here, which will continue to be the same so that we can include the files in our chroot:

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

Lines 225 to 227 in 9954e08

	typechecked_sources_request = Get(
	PythonSourceFiles, PythonSourceFilesRequest(partition.closure)
	)

What will change is this:

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

Lines 278 to 282 in 9954e08

	typechecked_srcs_snapshot = typechecked_sources.source_files.snapshot
	file_list_path = "__files.txt"
	python_files = "\n".join(
	determine_python_files(typechecked_sources.source_files.snapshot.files)
	)

We won't run on all of typechecked_sources anymore, only what was specified. (FYI, see https://www.pantsbuild.org/docs/rules-api-file-system#core-abstractions-digest-and-snapshot for what the "snaphsot" part means.)

"roots" vs. "closure"

Instead, we need to figure out which files were specified by the user. How? Looking back at the below definition, the field_set_addresses represent the roots, meaning what the user explicitly said to run on. Meanwhile, closure represents the transitive closure: all roots + transitive deps.

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

Lines 60 to 65 in 9954e08

	@dataclass(frozen=True)
	class MyPyPartition:
	field_set_addresses: FrozenOrderedSet[Address]
	closure: FrozenOrderedSet[Target]
	interpreter_constraints: PexInterpreterConstraints
	python_version_already_configured: bool

We want to operate on the roots, i.e. the field_set_addresses. We need to get the Python source files that belong specifically to those roots, whereas now we're getting it for the whole closure. We'll use the same approach as before, sayingGet(PythonSourceFiles, PythonSourceFilesRequest(..)), which will run a rule that gets us the relevant Python source files; only, we'll run it over just the roots, rather than the whole closure.

Part 1: refactor

First, there's a refactor to facilitate this change. Rather than storing on MyPyPartition the field field_set_addresses: FrozenOrderedSet[Address], change it to root_targets: FrozenOrderedSet[Target]. We're going to need Target objects in a moment, and this is also a clearer name.

Then, fix our rule that creates the partitions. Update to combined_roots.update(transitive_targets.roots):

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

Lines 370 to 372 in 9954e08

	for transitive_targets in all_transitive_targets:
	combined_roots.update(tgt.address for tgt in transitive_targets.roots)
	combined_closure.update(transitive_targets.closure)

You'll need to fix this too, use tgt.address for tgt in partitions.root_targets to convert from Target -> Address:

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

Line 243 in 9954e08

(addr for addr in partition.field_set_addresses),

Check that everything works by running ./pants typecheck src/python/pants/util/strutil.py.

Part 2: resolve the roots' files

See this:

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

Lines 222 to 227 in 9954e08

	plugin_sources_request = Get(
	PythonSourceFiles, PythonSourceFilesRequest(plugin_transitive_targets.closure)
	)
	typechecked_sources_request = Get(
	PythonSourceFiles, PythonSourceFilesRequest(partition.closure)
	)

First, that name typechecked_sources_request is confusing now. Really, it should probably be something like closure_sources_request. We're not going to typecheck every one of those files anymore.

Then, add something like roots_sources_request = Get(PythonSourceFiles, PythonSourceFilesRequest(partition.root_targets)). Feel free to use a better var name.

Add to this await MultiGet, which is what will cause the engine to actually execute the relevant code:

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

Lines 264 to 276 in 9954e08

	(
	plugin_sources,
	typechecked_sources,
	mypy_pex,
	requirements_pex,
	config_digest,
	) = await MultiGet(
	plugin_sources_request,
	typechecked_sources_request,
	mypy_pex_request,
	requirements_pex_request,
	config_digest_request,
	)

Finally, update this to point to your new PythonSourceFiles object with just the roots, rather than the original one with the closure.

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

Line 278 in 9954e08

typechecked_srcs_snapshot = typechecked_sources.source_files.snapshot

Update this:

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

Line 321 in 9954e08

description=f"Run MyPy on {pluralize(len(typechecked_srcs_snapshot.files), 'file')}.",

That should be it! Check that things are working how you want. It can be helpful to add logging statements to debug. https://www.pantsbuild.org/docs/rules-api-tips#tip-add-logging

--

This function will still stay the same to normalize file names, but the docstring should be tweaked as we won't run over every file anymore:

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

Lines 119 to 135 in 9954e08

	def determine_python_files(files: Iterable[str]) -> Tuple[str, ...]:
	"""We run over all .py and .pyi files, but .pyi files take precedence.
	MyPy will error if we say to run over the same module with both its .py and .pyi files, so we
	must be careful to only use the .pyi stub.
	"""
	result: OrderedSet[str] = OrderedSet()
	for f in files:
	if f.endswith(".pyi"):
	py_file = f[:-1] # That is, strip the `.pyi` suffix to be `.py`.
	result.discard(py_file)
	result.add(f)
	elif f.endswith(".py"):
	pyi_file = f + "i"
	if pyi_file not in result:
	result.add(f)
	return tuple(result)

Testing

You'll also want to check that mypy/rules_integration_test.py works correctly still. You may need to tweak some tests.

It'd be helpful to add a simple test that ensures this behavior works correctly. We're going to be changing this MyPy code in the future and we don't want to regress. Look at how the other tests are doing it to see how to set up, and also refer to https://www.pantsbuild.org/docs/rules-api-testing.

[Eric-Arellano]

Thanks to @hephex for fixing this!

[wimax-grapl]

Thanks for fixing this <3