Improving type support for `math.prod` #13572

cake-monotone · 2025-03-02T04:17:49Z

Related to #13501

The existing prod type stub did not properly support types like Fraction, complex, and Decimal provided by Python. To address this, I attempted to define the type stub in a way similar to builtins.sum.

Before

@overload
def prod(iterable: Iterable[SupportsIndex], /, *, start: SupportsIndex=1): ...
@overload
def prod(iterable: Iterable[SupportsFloatOrIndex], /, *, start: SupportsFloatOrIndex=1]): ...

Limitations

The previous type stub only checked for __index__ or __float__, but I believe this approach is entirely incorrect. There is no code in math_prod_impl that calls __index__ or __float__. As a result, the following cases fail:

import math

class Foo:
    def __index__(self) -> int:
        return 0

class Bar:
    def __float__(self) -> float:
        return 0.0

math.prod([Foo(), Foo()])  # Fails
math.prod([Bar(), Bar()])  # Fails

Additionally, the old stub did not account for empty iterables:

float_list: list[float] = []
math.prod(float_list)  # Returns 1 at runtime, but inferred as int.

After

_MultiplicableT1 = TypeVar("_MultiplicableT1", bound=SupportsMul[Any, Any])
_MultiplicableT2 = TypeVar("_MultiplicableT2", bound=SupportsMul[Any, Any])

class _SupportsProdWithNoDefaultGiven(SupportsMul[Any, Any], SupportsRMul[int, Any], Protocol): ...
_SupportsProdNoDefaultT = TypeVar("_SupportsProdNoDefaultT", bound=_SupportsProdWithNoDefaultGiven)

@overload
def prod(iterable: Iterable[bool | _LiteralInteger], /, *, start: int = 1) -> int: ...  # type: ignore[overload-overlap]
@overload
def prod(iterable: Iterable[_SupportsProdNoDefaultT], /) -> _SupportsProdNoDefaultT | Literal[1]: ...
@overload
def prod(iterable: Iterable[_MultiplicableT1], /, *, start: _MultiplicableT2) -> _MultiplicableT1 | _MultiplicableT2: ...

This implementation is largely derived from sum, so it shares the same limitations. You can find related PRs for sum here: #7578 and #8000.

Since those PRs were created in 2022, I tried different ways to improve this using newer typing features introduced since then, but none worked out. If anyone has a better idea, I'd love to hear it!

Improvements

With this change, prod correctly handles cases like:

prod([complex(1, 2), complex(3, 4)])  # Returns a complex number!
prod([Fraction(1, 2), Fraction(1, 2)])  # Returns a Fraction!

Limitations

This implementation still produces many false positives. Non-numeric types in Python use the multiplication operator in many different ways due to syntactic sugar, which makes precise typing pretty difficult.

Expression	Runtime Result	Actual Type	Expected Type
`prod("abcde")`	Runtime Error		`str \| Literal[1]`
`prod([1, 2, 3], start="a")`	"aaaaaa"	`str`	`int \| str`
`prod([2], start=(1, 2))`	`(1, 2, 1, 2)`	`tuple`	`int \| tuple[Literal[1], Literal[2]]`

github-actions · 2025-03-02T04:28:53Z

According to mypy_primer, this change has no effect on the checked open source code. 🤖🎉

srittau

Thank you, also for the extensive tests.

cake-monotone added 2 commits March 2, 2025 13:06

Improving type support for math.prod

bde6425

fix pr number in the comment

a57ec2e

srittau approved these changes Mar 3, 2025

View reviewed changes

srittau merged commit 9f11db4 into python:main Mar 3, 2025
55 checks passed

mmingyu pushed a commit to mmingyu/typeshed that referenced this pull request May 16, 2025

Improving type support for math.prod (python#13572)

0e95274

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Improving type support for `math.prod` #13572

Improving type support for `math.prod` #13572

Improving type support for math.prod #13572

Improving type support for math.prod #13572

Conversation

Before

Limitations

After

Improvements

Limitations

Choose a reason for hiding this comment

Improving type support for `math.prod` #13572

Improving type support for `math.prod` #13572