[FIXED] Implementing operators using data model objects and checking with pyright

Issue

I am working on a class where I have to implement arithmetic operators (+, -, *, /), where each operator has multiple overloads, but the overloads are all the same. Instead of copying the overloads to each operator, I thought I would implement it as a data model object, something along the lines of this:

from typing import Callable as Fn, Any, overload
import operator


class Apply:
    """Apply an operator to an object."""

    def __init__(self, op: Fn[[Any, Any], Any], obj: Any) -> None:
        self.op = op
        self.obj = obj

    # Two mock overloads...
    @overload
    def __call__(self, x: int) -> str:  ...
    @overload
    def __call__(self, x: str) -> int:  ...

    def __call__(self, x: int | str) -> str | int:
        ...


class Op:
    """Data model object for an operator."""

    def __init__(self, op: Fn[[Any, Any], Any]) -> None:
        self.op = op

    def __get__(self, obj: Any, _: Any) -> Apply:
        return Apply(self.op, obj)


class Foo:
    __add__ = Op(operator.add)
    __mul__ = Op(operator.mul)


foo = Foo()
a: str = foo.__add__(2)    # works fine
b: int = foo.__mul__("2")  # works fine
_ = foo + 1    # type error
_ = foo * "2"  # type error

The idea would be to have all the operator overloads implemented once, with a dispatch to an operator for the dynamic behaviour. That way, I avoid multiple copies of the boilerplate code, and I still get all my operators with the overload type annotations.

But the last two lines give me type errors when I check with pyright (playground). It seems to be working with mypy, but I use pylance+pyright for my project and would rather not change type checker.

Is there a way to get this to work with pyright?

Solution

It looks like pyright just needs a bit of extra help (see Pyright playground):

class Op:
    """Data model object for an operator."""

    def __init__(self, op: Fn[[Any, Any], Any]) -> None:
        self.op = op

    def __get__(self, obj: Any, _: Any) -> Apply:
        return Apply(self.op, obj)
    
    __call__: Apply  # Helper annotation

>>> foo = Foo()
>>> reveal_type(foo.__add__(2))    # str
>>> reveal_type(foo.__mul__("2"))  # int
>>> reveal_type(foo + 1)           # str
>>> reveal_type(foo + "2")         # int

Answered By - dROOOze