Functions aren't always able to match TypeVar to method argument types

[smeznaric]

Environment data

• Language Server version: 2021.4.2
• OS and version: CentOS Linux release 7.5.1804
• Python version (& distribution if applicable, e.g. Anaconda): 3.6.3

Code Snippet / Additional information

from typing import TypeVar
import pandas as pd

DfSer = TypeVar("DfSer", pd.Series, pd.DataFrame)


def divide_them(pd1: DfSer, pd2: DfSer) -> DfSer:
    return pd1.divide(pd2)


def divide_them_2(pd1: DfSer, pd2: DfSer) -> DfSer:
    if isinstance(pd1, pd.Series):
        return pd1.divide(pd2)
    else:
        assert isinstance(pd1, pd.DataFrame)
        return pd1.divide(pd2)


def divide_them_3(pd1: DfSer, pd2: DfSer) -> DfSer:
    if isinstance(pd2, pd.Series):
        return pd1.divide(pd2)
    else:
        assert isinstance(pd2, pd.DataFrame)
        return pd1.divide(pd2)


def divide_them_4(pd1: DfSer, pd2: DfSer) -> DfSer:
    if isinstance(pd1, pd.Series) and isinstance(pd2, pd.Series):
        return pd1.divide(pd2)
    else:
        assert isinstance(pd1, pd.DataFrame) and isinstance(pd2, pd.DataFrame)
        return pd1.divide(pd2)

Expected behaviour

Nothing should be underlined.

Both DataFrame and Series have a divide method. Series method accepts a Series and DataFrame method accepts either a DataFrame or a Series.

Actual behaviour

In the divide_them function we get for argument pd2:
Argument of type "DfSer@divide_them" cannot be assigned to parameter "other" of type "int | float | ndarray | Dict[_str, ndarray] | Sequence[Unknown] | Index[Unknown] | Series[Unknown]" in function "divide"

Similar issues in other functions.

We avoid errors if we explicitly check for both arguments like in divide_them_4.

[erictraut]

Thanks for the bug report. This bug looks similar to as this one: #1180. It is very difficult to repro in pyright, but it appears readily in pylance. I'll look into it in more detail.

[erictraut]

I'm pretty sure I've fixed this bug in the pyright source base.

@jakebailey, we should double check that the problem no longer appears in pylance when you pull the code over and integrate next time.

[jakebailey]

I'm not so sure; in my pull I see this:

[jakebailey]

Most recent commit didn't help, unfortunately.

[erictraut]

I'm able to repro the core problem in this bug report. Here's a simplified repro, thanks to @jakebailey:

from typing import TypeVar

class A:
    def method(self, x: "A") -> "A": ...

class B:
    def method(self, x: "B") -> "B": ...

T = TypeVar("T", A, B)

def check(x: T, y: T) -> T:
    return x.method(y)

Supporting this use case is going to be a challenge. It requires that constraints imposed by a constrained TypeVar need to be distributed across a function call expressions when evaluating the arguments for that call. In the example above, for example, the type of argument expression y needs to be constrained to either A or B depending on the type of x. The type checking engine already supports such constraints when used with other expression types like binary operators (such as "+" or "*"), but it doesn't currently support constraints for arbitrary call expressions. It does have some support for isinstance type narrowing based on these constraints, which is why some of the attempts to use isinstance work in the example above.

[erictraut]

I've added the support for propagating constraints to call arguments when the call portion of a call expression imposes such a constraint. For example, the sample above uses x.method as the call expression, and it imposes constraints on the arguments based on the fact that x is a constrained TypeVar. That means when evaluating the argument expressions (like y) we need to apply the constraint (in this case either limiting it to A or B).

The fix will be included in the next release of pylance.

[smeznaric]

Brilliant, thanks!

[jakebailey]

This issue has been fixed in version 2021.4.3, which we've just released. You can find the changelog here: https://github.com/microsoft/pylance-release/blob/main/CHANGELOG.md#202143-29-april-2021