Skip to content

Improve precision of itertools.product types when called with 7 or more arguments#11163

Merged
JelleZijlstra merged 1 commit intopython:mainfrom
winestone:feature/better-itertools-product-typing
Feb 16, 2024
Merged

Improve precision of itertools.product types when called with 7 or more arguments#11163
JelleZijlstra merged 1 commit intopython:mainfrom
winestone:feature/better-itertools-product-typing

Conversation

@winestone
Copy link
Copy Markdown
Contributor

@winestone winestone commented Dec 14, 2023

Previously when the following expression would yield a type with Any:

args: List[List[int]]
itertools.product(*args)  # product[tuple[Any, ...]]

but this is ok:

itertools.product(*args, repeat=1)  # product[tuple[int, ...]]

This PR attempts to fix that but do let me know if I've got anything wrong (I only did some local testing in pyright)!

EDIT: Fixed a typo args: List[int] -> args: List[List[int]] found by @hauntsaninja #11163 (review) , thanks!

@github-actions
Copy link
Copy Markdown
Contributor

Diff from mypy_primer, showing the effect of this PR on open source code:

jax (https://github.com/google/jax)
+ jax/_src/scipy/ndimage.py:123: error: Argument 1 to "_nonempty_prod" has incompatible type "tuple[Array | Any | int | float | complex, ...]"; expected "Sequence[Array]"  [arg-type]

Copy link
Copy Markdown
Collaborator

@hauntsaninja hauntsaninja left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks, this looks good to me (although I think you mean list[list[int]] in PR description). The mypy_primer hit is unfortunate, but in general it's hard to preserve types through zip(*list) like constructs

@Daverball
Copy link
Copy Markdown
Contributor

While this is probably not a big deal, since a product of more than six iterables will rarely come up, this does mean that anything using more than that will now will most likely result in a product[tuple[object, ...]], unless the iterables are already uniform, which will be more annoying to deal with than a product[tuple[Any, ...]].

My intuition initially was that mypy would reject mixing non-uniform iterables for the final overload, but since Iterable is covariant it will implicitly cast all the iterables to the common base type for the final overload, rather than rejecting them based on non-uniformity.

It's probably still a good trade-off to support the nested iterable case, since it is a more common use-case, just thought I'd mention it.

@JelleZijlstra JelleZijlstra merged commit f3c7c48 into python:main Feb 16, 2024
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants