How to write type annotation when the returned value is a particular subclass

89 Views Asked by At

I have several subclasses of the same class. I implemented is_instance methods and want to use them to detect the type of the object, but Pylance says that the superclass is incompatible with the subclass. This is the simlified example:

class Employee:

    def is_worker(self) -> bool:
        return False

    def is_manager(self) -> bool:
        return False


class Worker(Employee):

    def is_worker(self) -> bool:
        return True


class Manager(Employee):

    def is_manager(self) -> bool:
        return True


def get_employees_queue() -> list[Employee]:
    # Simplified example just to give a clue about possible types
    return [Manager(), Worker()]


def get_a_worker() -> Worker:
    employees = get_employees_queue()
    for employee in employees:
        if employee.is_worker():
            return employee    # here Pylance says: Expression of type "Employee"
                               # cannot be assigned to return type "Worker"
                               # "Employee" is incompatible with "Worker"
    raise RuntimeError('No workers found')

Tried solutions:

  1. Built in isinstance - solves the Pylance issue, but I want to use before mentioned methods here and in other modules since they are more readable.
  2. Changing returned type from Worker to Emploee as well helps, but why should I do it if I know that there always be Worker?
  3. Add employee: Worker before the return statement. In that case Pylance shows the error on line for employee in employees:.

So the question is how to say to Pylance that the returned type is always Worker?

Pyright Playground

1

There are 1 best solutions below

3
On BEST ANSWER

Simplest solution here is probably to use isinstance() instead of your custom type check methods. I tried to make it work using overload and TypeGuard features, but that was a no go. It appears TypeGuard's cannot be used for instance methods.

So, for simplest solution, rewrite the above.

  1. Get rid of functions 'is_worker' and 'is_manager' from 'Employee' model
  2. Either rename those functions to def __instancecheck__(self) -> bool or delete them altogether from the subclasses too.

The __instancecheck__ dunder method allows you to customize the logic for checking whether something is an instance. If you have special needs for checking, you can put them there. From your example, you probably have no special requirements, so you can just drop those methods completely.

  1. Change your code to something like this:
    for employee in employees:
        if isinstance(employee, Worker):
            return employee    
    raise RuntimeError('No workers found')

Doing it this way should make your type-checking system happy and get rid of the warning you are seeing from Pylance (I don't have Pylance, I can't confirm this).

Short answer for most use cases: don't use is_XYZ(self) methods. Just use Python's existing isinstance() feature.

Solution using TypeGuard

You can see it in action here. Note that this solution requires 3.12 and does not (cannot) use instance methods

from typing import Any, TypeGuard


def is_worker(obj: Any) -> TypeGuard["Worker"]:
    return isinstance(obj, Worker)

def is_manager(obj: Any) -> bool:
    return False

class Employee:
    pass

class Worker(Employee):
    pass

class Manager(Employee):
    pass
    

def get_employees_queue() -> list[Employee]:
    # Simplified example just to give a clue about possible types
    return [Manager(), Worker()]


def get_a_worker() -> Worker:
    employees = get_employees_queue()
    for employee in employees:
        if is_worker(employee):
            return employee
    raise RuntimeError('No workers found')