Typing a function to differentiate classes based on an attribute

Question

I am currently trying to add type hints for a library that exposes to the end user a function like the following (simplified for clarity):

def bar(obj):
    if hasattr(obj, "foo"):
        return obj.foo()
    return obj

Because the obj parameter can be defined by the end user, I need to come-up with a generic way to convey intent and make sure that the bar function gives a precise type as return value:

if the argument has an foo attribute that can be callable without any parameters, returns the return type of that attribute;
else returns the type of obj.

My best try so far is the following:

from typing import Protocol, TypeVar, Union, cast

T = TypeVar("T")
U = TypeVar("U", covariant=True)

class HasFoo(Protocol[U]):
    def foo(self) -> U:
        ...

def bar(obj: Union[HasFoo[U], T]) -> Union[U, T]:
    if hasattr(obj, "foo"):
        return cast(HasFoo[U], obj).foo()
    return cast(T, obj)

It seems to work properly for classes that don't have a foo attribute:

from typing_extensions import assert_type

class A:
    pass

a = A()

bar_a = bar(a)
assert_type(bar_a, A)
assert bar_a == a

However it does not work for classes that do have the foo attribute (mypy errors in comments):

class B:
    def foo(self) -> int:
        return 42

b = B()

bar_b = bar(b)
# ^ error: Argument 1 to "bar" has incompatible type "B"; expected "HasFoo[]"  [arg-type]
# ^ error: Need type annotation for "bar_b"  [var-annotated]

assert_type(bar_b, int)
# ^ error: Expression is of type "Any", not "int"

assert bar_b == 42

The error I'm not understanding is Argument 1 to "bar" has incompatible type "B"; expected "HasFoo[]" [arg-type].

The issue might be coming from the fact that for the Union[HasFoo[U], T] parameter, mypy does not know how to distinct one from the other.

I have noticed that if I create the following function instead, everything works as expected:

from typing import Iterable


def bar2(obj: Union[HasFoo[U], Iterable[T]]) -> Iterable[Union[U, T]]:
    if hasattr(obj, "foo"):
        yield cast(HasFoo[U], obj).foo()
    else:
        yield from cast(Iterable[T], obj)


bar2_a = bar2([a])
assert_type(bar2_a, Iterable[A])
assert list(bar2_a) == [a]


bar2_b = bar2(b)
assert_type(bar2_b, Iterable[int])
assert list(bar2_b) == [42]

I suspect that this is because this time there is no common ground between HasFoo[U] and Iterable[T].

With all that being said, how can I properly type the initial bar function?

I am using python 3.10.6 and mypy 0.971

user_9232196 · Answer

Using @overload makes this relatively painless - just type the two different possibilities as two overloads, then leave the hints off the implementation.

from typing import Protocol, TypeVar, overload

T = TypeVar("T")
U = TypeVar("U", covariant=True)

class HasFoo(Protocol[U]):
    def foo(self) -> U:
        ...

@overload
def bar(obj: HasFoo[U]) -> U:
    ...

@overload
def bar(obj: T) -> T:
    ...

def bar(obj):
    if hasattr(obj, "foo"):
        return obj.foo()
    return obj

Using reveal_type with your examples yields the expected results: bar_a is an A and bar_b is an int.

Typing a function to differentiate classes based on an attribute

1 Answers