Regression with Overloads and lambda

[hmc-cs-mdrissi]

Describe the bug I'm unsure on exact issue. I know overload/lambda are key but still unsure what's special about this example vs similar ones I have. I've been able to minimize it down to ~60 line program and have stubbed out implementation.

from __future__ import annotations

from typing import Callable, Optional, TypeVar, Union, overload

_ExpectedT = TypeVar("_ExpectedT")

class Tensor:
    ...

class RaggedTensor:
    ...

class SparseTensor:
    ...

class RegisteredFeature:
    ...

class FeatureContainer:
    def get(self, feature: FeatureRef, t: type[_ExpectedT]) -> _ExpectedT:
        ...

FeatureRef = Union[Callable[[FeatureContainer], object], RegisteredFeature]

def densify(
    sparse_feature: Tensor | SparseTensor,
) -> Tensor:
    ...

@overload
def as_string(
    input: Tensor,
) -> Tensor:
    ...

@overload
def as_string(
    input: RaggedTensor,
) -> RaggedTensor:
    ...

def as_string(
    input: RaggedTensor | Tensor,
) -> RaggedTensor | Tensor:
    ...

def feauture_fn(dropoff_census_tract: RegisteredFeature) -> None:
    _: FeatureRef = lambda x: as_string(densify(x.get(dropoff_census_tract, Union[Tensor, SparseTensor])))

One small variation that passes type checking is to simplify as_string definition to,

def as_string(
    input: Tensor,
) -> Tensor:
    ...

By removing overloaded definition it passes. To Reproduce Run pyright 1.1.305. Produces false positive error on last line of,

Argument of type "Tensor" cannot be assigned to parameter "input" of type "RaggedTensor" in function "as_string"
  "Tensor" is incompatible with "RaggedTensor"

Except as_string has one overload that does accept a Tensor as input.

Expected behavior No error. This code passed type checking on 1.1.303. It started failing on 1.1.304. In particular as_string function has two overloads but if you delete second overload and make it no longer overloaded it passes type checking. I wouldn't expect adding an overload to cause it to not be able to match first function anymore.

For lambda it's related to bidirectional inference. If I remove : FeatureRef and make it just ` = lambda x: as_string(densify(x.get(dropoff_census_tract, Union[Tensor, SparseTensor])))` the error disappears. In real usage the bidirectional inference happens because I pass lambda to another function that is annotated for it's arguments.

edit: I have a second somewhat similar example that I'm thinking is same bug. This re-uses some of definitions from first example. Similar to first example it passes type checking <=1.1.303 and fails to check >1.1.304.

@overload
def apply_binary_op(
    binary_op: Callable[[Tensor, Tensor], Tensor],
    left_operand: Tensor,
    right_operand: Tensor,
) -> Tensor:
    ...

@overload
def apply_binary_op(
    binary_op: Callable[[Tensor, Tensor], Tensor],
    left_operand: SparseTensor,
    right_operand: Tensor,
) -> SparseTensor:
    ...

@overload
def apply_binary_op(
    binary_op: Callable[[Tensor, Tensor], Tensor],
    left_operand: Tensor,
    right_operand: SparseTensor,
) -> SparseTensor:
    ...

def apply_binary_op(
    binary_op: Callable[[Tensor, Tensor], Tensor],
    left_operand: Union[Tensor, SparseTensor],
    right_operand: Union[Tensor, SparseTensor],
) -> Union[Tensor, SparseTensor]:
    ...

def greater(x: Tensor, y: Tensor) -> Tensor:
    ...

def _foo(container: FeatureContainer, tips: RegisteredFeature, threshold_tip: Tensor) -> None:
    apply_binary_op(greater, densify(container.get(tips, Union[Tensor, SparseTensor])), threshold_tip)

[erictraut]

This issue is related to this one. It has to do your use of a union (Union[Tensor, SparseTensor]) in a value expression. Unions should generally be used only within type expressions. They are officially supported in a few special cases such as TypeVar bounds, constraints, and arguments passed to isinstance or issubclass. You're using it in other places where they are not supported by a strict reading of PEP 484.

As we discussed in the context of the other bug report, the type system doesn't currently support TypeForm, so I'm inclined to support your rather unorthodox use here. I've made a change that fixes this issue. It will be included in the next release.

[erictraut]

This is addressed in pyright 1.1.306, which I just published. It will also be included in a future release of pylance.