Does pyright treat intermediate object types differently to objects created in a return statement?

[danbraunai-apollo]

Just ran into this on torch==2.5.1 (notably, not on torch==2.3.1).

import torch

class A:
    def __init__(self):
        # Note that torch.nn.Parameter is a child of torch.Tensor
        self.a = torch.nn.Parameter(torch.empty(2))

    def get_param_a_1(self) -> dict[str, torch.Tensor]:
        param_a = {"a": self.a}
        # Raises Type "dict[str, Parameter]" is not assignable to return type "dict[str, Tensor]"
        return param_a

    def get_param_a_2(self) -> dict[str, torch.Tensor]:
        return {"a": self.a}  # No pyright error

    def get_param_a_3(self) -> dict[str, torch.Tensor]:
        param_a: dict[str, torch.Tensor] = {"a": self.a}
        return param_a  # No pyright error

This was an unexpected gotchya for me. Is this behaviour intended? Is there a good reason that pyright doesn't type objects that are created in a return statement the same way it types objects assigned to a variable? Thanks!

[erictraut]

Yes, this is expected behavior. The expression within a return statement is evaluated using bidirectional type inference. The return type provides the "expected type" as context. When the same expression is assigned to a variable whose type is not declared, regular inference rules apply, since there is no context.

If you want get_param_a_1 to work the same way as get_param_a_2, you can provide an explicit type declaration for the param_a variable, as you've done in get_param_a_3.