Add support for gRPC mypy stub files

[jyggen]

Problem

Pants currently supports generating mypy stubs for the generated protobuf code, which makes mypy very happy. However, when type checking a codebase that has generated gRPC code as well, mypy throws quite the hissy fit due to the lack of gRPC stubs.

Solution

mypy-protobuf can (since version 2.0 at least) generate gRPC stubs as well, and it can be used in Pants with some minor changes. It will, of course, only work on mypy-protobuf 2.0+, and anyone using an older version will have a bad time. I'm not familiar enough with the Pants codebase to know if adding a version check to if downloaded_grpc_plugin: will be enough to prevent failure or if VenvPexRequest will crash and burn if a potentially invalid binary is referenced.

Result

Stubs are now generated for the generated gRPC code as well! Awesome!

[Eric-Arellano]

This is awesome! Thank you so much @jyggen!

It will, of course, only work on mypy-protobuf 2.0+, and anyone using an older version will have a bad time. I'm not familiar enough with the Pants codebase to know if adding a version check to if downloaded_grpc_plugin: will be enough to prevent failure or if VenvPexRequest will crash and burn if a potentially invalid binary is referenced.

Hm, yeah, tricky. We would need to check if mypy-protobuf is 2+, but that's non-trivial because the requirement string could be something like mypy-protobuf>=1.0. I don't know of a library to do this. We have code like this for Python versions:

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

Lines 223 to 230 in f411f0d

	def _includes_version(self, major_minor: str, last_patch: int) -> bool:
	patch_versions = list(reversed(range(0, last_patch + 1)))
	for req in self:
	if any(
	req.specifier.contains(f"{major_minor}.{p}") for p in patch_versions # type: ignore[attr-defined]
	):
	return True
	return False

But that's not very performant to iterate through every possible version and it's more complexity than we ideally want.

I think the intersection is small enough here that it's okay to break mypy-protobuf 1.x users; to be hit by this, you must have:

1. enabled [python-protobuf].mypy_plugin
2. pinned [python-protobuf].mypy_plugin_version to 1.x
3. be using grpc=True in some protobuf_library targets

@stuhood @benjyw what do you think about the backwards compat? The concerning part is that there is no remedy if you're in that small audience other than bumping to 2.x for the plugin..we permanently break gRPC users who must stay on 1.x of the plugin.

--

@jyggen, would you be willing to please break out a precursor PR that simply bumps --mypy-plugin-version to 2.4? It would be great for the changelog to have an entry "Bump default for [python-protobuf].mypy_plugin_version to 2.4", given that they made API changes.

[benjyw]

Thanks for the contribution @jyggen! Yeah the backwards compat question is tricky.

I think we can mitigate by capturing the semver from the requirement string and just checking the major version. In the case of mypy-protobuf>=1.0 we'd just be stricter than we need to be. But we should probably also issue a deprecation warning if we detect you're on 1.x, and remove support for it after the cycle.

[jyggen]

@Eric-Arellano Sure! Done in #11662.

[Eric-Arellano]

I think we can mitigate by capturing the semver from the requirement string and just checking the major version

Yeah I think you're right that we should do this. And it's less painful than I expected, there are few 1.x releases for the plugin and they don't use patch versions, so we don't need to check a ton of versions to see if 1.x is in use: https://pypi.org/project/mypy-protobuf/#history

@jyggen could you please add a check if any of 1.0 up to 1.27 can be matched by --python-protobuf-mypy-plugin-version? Their releases only go up to 1.24, but better to be cautious if they do any 1.x releases in the future.

See this example:

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

Lines 223 to 230 in f411f0d

	def _includes_version(self, major_minor: str, last_patch: int) -> bool:
	patch_versions = list(reversed(range(0, last_patch + 1)))
	for req in self:
	if any(
	req.specifier.contains(f"{major_minor}.{p}") for p in patch_versions # type: ignore[attr-defined]
	):
	return True
	return False

That is, first use Requirement.parse() from pkg_resources, then iterate over 1.{minor} for the range 0-28. Two notes on this:

1. We're checking "is this 1.x?", rather than "is this 2.x?" because that's more future-proof. We don't know what the future release versions will be, e.g. a 3.x release.
2. We are overly conservative here, which could be the wrong decision..This will not match on a loose constraint like mypy-protobuf, as technically it does match 1.x even though it also matches 2.x and would likely use 2.x (unless another requirement or a constraints file pin it to 1.x). You'd need to use something more constrained like mypy-protobuf>=2.0 or mypy-protobuf==2.4 to get your new feature working. See https://github.com/pantsbuild/pants/pull/11585/files#diff-f1be1453b60d548e4f74bed54e90eaa2e6bd25bcdcca107323e5e6e5d8ddac93R66-R90 for an alternative approach. What do you think? I'm leaning towards us being more conservative because this feature will break 1.x users, with no remedy. But it's also confusing to a new user why mypy-protobuf doesn't generate gRPC.

It probably makes sense to write that as a top-level function so that you can write a unit test, e.g.

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

Lines 140 to 152 in e3b3136

	@pytest.mark.parametrize(
	"constraints",
	[
	["CPython>=2.7,<3"],
	["CPython>=2.7,<3", "CPython>=3.6"],
	["CPython>=2.7.13"],
	["CPython>=2.7.13,<2.7.16"],
	["CPython>=2.7.13,!=2.7.16"],
	["PyPy>=2.7,<3"],
	],
	)
	def test_interpreter_constraints_includes_python2(constraints) -> None:
	assert PexInterpreterConstraints(constraints).includes_python2() is True

[jsirois]

Keep in mind it's simple to just check the actual version. 1st create a VenvPex with the 1script, then run PEX_TOOLS=1 info against it to get its PEX-INFO which will have a distrubutions dict of the actual resolved wheel hashes / names from which you can parse the exact version. If that version is high enough, fire off a 2nd VenvPex request to create with 2 scripts, using 1st VenvPex as the main source via PEX_PATH.

[Eric-Arellano]

Keep in mind it's simple to just check the actual version.

True. @jyggen I'll put up a precursor PR to add a helper rule to facilitate this. Scratch my above recommendation about doing static analysis on the requirement string.

[jsirois]

Sorry to armchair, but I'll be AFK here in a few minutes. I wanted to make sure you all were also aware the parser needed for this is already in our dependency set: https://github.com/pypa/packaging/blob/53fd698b1620aca027324001bf53c8ffda0c17d1/packaging/utils.py#L88-L89

[jyggen]

So what's the next step? Add an additional check so we only append the new flags when 2.x of the plugin is used ? I took a quick look at it earlier but couldn't find any (to me) obvious way to get the version of the plugin that ends up being installed, but if someone can point me in the right direction I could take a proper look at it.

[Eric-Arellano]

So what's the next step?

Thanks for checking in. John is about to land #11684, which will unblock me from landing #11665 so that we have a way to get information from what was actually used in the PEX, rather than relying on static string analysis. Then, you'll use that in this PR to implement your logic.

I'll send some further thoughts once I land #11665, I'm still trying to see what the final API will look like.

[benjyw]

@jyggen thanks for your patience as we work this through. Should get it in within a few days!

[jyggen]

@benjyw No worries at all! Just wanted to check if the ball was in my court or not :)

[Eric-Arellano]

Okay, #11665 has landed:

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

Lines 1018 to 1031 in 1902297

	@dataclass(frozen=True)
	class PexDistributionInfo:
	"""Information about an individual distribution in a PEX file, as reported by `PEX_TOOLS=1
	repository info -v`."""
	project_name: str
	version: packaging.version.Version
	requires_python: packaging.specifiers.SpecifierSet | None
	requires_dists: tuple[Requirement, ...]
	class PexResolveInfo(Collection[PexDistributionInfo]):
	"""Information about all distributions resolved in a PEX file, as reported by `PEX_TOOLS=1
	repository info -v`."""

Start by reverting setting protoc_gen_mypy_grpc_script in the original mypy_pex await Get.

In that if python_protobuf_subsystem.mypy_plugin block, add a nested if request.protocol_target.get(ProtobufGrpcToggle).value to see if gRPC is being used. (You might want to factor this out into a dedicated variable because it will get used later when deciding to download gRPC or not.) In the block:

1. Add await Get(PexResolveInfo, VenvPex, mypy_pex). Then, look if there's an entry with project_name = mypy-protobuf and then probably version >= 2.0 (see https://packaging.pypa.io/en/latest/version.html). Note that the user can override --python-protobuf-mypy-plugin-version to whatever arbitrary string they want, so it's not guaranteed the pex will actually have mypy-protobuf resolved.
2. If the above conditions are met, override mypy_pex with await Get(VenvPex, VenvPexRequest), with all the same kwargs as before except both bin_names this time.
   • John had recommended using pex_path, which merges PEXes together at runtime, but that's not yet possible with our VenvPex abstraction and we don't need to block this feature on that. I can fix that in a followup PR. We'll have wasted work this way, that you resolve 2 similar PEXes twice, but that gets cached to disk at least and will be fixed.
3. Update the argv iff the conditions from step 1 were met.

Then, final step is to update the new test you added to check that if --mypy-plugin-version is 1.x, we do not generate gRPC and things still succeed. You will need to set this via rule_runner.set_options():

pants/src/python/pants/backend/codegen/protobuf/python/rules_integration_test.py

Lines 43 to 52 in 0f27e41

	options = [
	"--backend-packages=pants.backend.codegen.protobuf.python",
	f"--source-root-patterns={repr(source_roots)}",
	]
	if mypy:
	options.append("--python-protobuf-mypy-plugin")
	rule_runner.set_options(
	options,
	env_inherit={"PATH", "PYENV_ROOT", "HOME"},
	)

[benjyw]

Thanks! @Eric-Arellano , sounds like you'll follow up on that TODO and a couple of others?

[Eric-Arellano]

@Eric-Arellano , sounds like you'll follow up on that TODO and a couple of others?

Thank you! Yes, I will get --pex-path working to avoid the unnecessary Pex resolve, and will also address the two places I see code duplication if Jonas doesn't beat me to it.

This looks great! Thanks for the patch and for iterating on this.

[jyggen]

Thanks! The code duplication should be fixed. While I was at it I also updated the TODO comment to be a tad more accurate.. hopefully!

[jsirois]

Thanks!

I've got a change here that will allow the pex_path cleanup in another way: #11693