Mypy: looking for the perfect signature for an average function

user_1973926 0
Viewed 418

I'm trying to come up with the perfect function signature for the following function (Python 3.6, mypy 0.521):

def avg(xs):
    it = iter(xs)
    try:
        s = next(it)
        i = 1
    except StopIteration:
        raise ValueError("Cannot average empty sequence")
    for x in it:
        s += x
        i += 1
    return s / i

The nice thing about this code, is that it works with, and produces the correct result for iterables of int, float, complex, but also for datetime.timedelta. Problems pop up when trying to add a signature. I've tried the following:

def avg(xs: t.Iterable[t.Any]) -> t.Any: ...

But now, caller needs to cast the result.

def avg(xs: t.Iterable[T]) -> T: ...

This fails because T does not support addition nor division.

N = TypeVar("N", int, float, complex, datetime.timedelta)
def avg(xs: t.Iterable[N]) -> N: ...

Fails because int / int is a float; using // gives the wrong result for almost everything else. Also sucks, because the code should work for other types, as long as addition and division are supported.

N = TypeVar("N", float, complex, datetime.timedelta)
def avg(xs: t.Iterable[N]) -> N: ...

This is almost perfect, but again, if someone later decides to throw quaternions at it, mypy will complain.

...then I was also trying something with abc and typing.overload but that got me nowhere.

What would be the most elegant solution that would pass under mypy --strict?

mypy python
1 Answers

So, unfortunately, the numerics system in Python/PEP 484 is currently a bit of a mess.

We technically have a "numeric tower" that's supposed to represent a set of ABCs that all "number-like" entities in Python are supposed to obey.

Furthermore, many of the built-in types in Python (such as int, float, complex, and timedelta) don't inherit from these ABCs in typeshed -- this means that those ABCs are basically unusable (except in cases where you define custom types that explicitly inherit from those ABCs).

And to compound the problem, the numbers module is largely dynamically typed in typeshed -- I took a swing at fixing the numbers module about a year or so ago, and my recollection was that mypy at the time wasn't powerful enough to accurately type the numeric tower.

That situation might be fixed today, but that's all more or less moot, because mypy recently implemented experimental support for Protocols (e.g. structural typing)! It turns out this is exactly what we need to address your problem and eventually fix the numeric tower (once Protocols are added to PEP 484 and the typing module).

For now, what you need to do is:

  1. Install the typing_extensions module (python3 -m pip install typing_extensions)
  2. Install the latest version of mypy from Github (run python3 -m pip install -U git+git://github.com/python/mypy.git)

We can then define a protocol for a "supports add or divide" type like so:

from datetime import timedelta

from typing import TypeVar, Iterable
from typing_extensions import Protocol

T = TypeVar('T')
S = TypeVar('S', covariant=True)

class SupportsAddAndDivide(Protocol[S]):
    def __add__(self: T, other: T) -> T: ...

    def __truediv__(self, other: int) -> S: ...

def avg(xs: Iterable[SupportsAddAndDivide[S]]) -> S:
    it = iter(xs)
    try:
        s = next(it)
        i = 1
    except StopIteration:
        raise ValueError("Cannot average empty sequence")
    for x in it:
        s += x
        i += 1
    return s / i

reveal_type(avg([1, 2, 3]))
reveal_type(avg([3.24, 4.22, 5.33]))
reveal_type(avg([3 + 2j, 3j]))
reveal_type(avg([timedelta(1), timedelta(2), timedelta(3)]))

Running this using mypy produces the following output, as desired:

test.py:27: error: Revealed type is 'builtins.float*'
test.py:28: error: Revealed type is 'builtins.float*'
test.py:29: error: Revealed type is 'builtins.complex*'
test.py:30: error: Revealed type is 'datetime.timedelta*'
user_6465430
Related